mihir2007/Cyber-Risk
0
1"""2db_init.py3----------4Database initialization and bootstrapping script for the CRQ platform.5 6Capabilities:71. Tests the PostgreSQL connection and retrieves server telemetry.82. Drops and recreates all database tables with safety confirmation guards.93. Populates initial seed data:10 - Enterprise IT/OT assets spanning Critical, Medium, and Low tiers.11 - Authentic CVE records with CVSS 3.1 and EPSS distributions.12 - Comprehensive security controls catalogue.13 - Enterprise NetworkEdge topology for dynamic lateral-movement analysis.144. Verifies attack graph connectivity and Dijkstra reachability from the perimeter.15 16Usage:17 python db_init.py # Prompts before dropping tables18 python db_init.py --force # Non-interactive automated deployment19 python db_init.py --test-only # Verifies DB connectivity without changes20"""21 22from __future__ import annotations23 24import argparse25import sys26from typing import Any27 28from sqlalchemy import text29from sqlalchemy.exc import OperationalError, SQLAlchemyError30from sqlalchemy.orm import Session31 32from database import Base, SessionLocal, engine33from graph_engine import AttackGraphEngine34from models import Asset, DataLineageFlow, NetworkEdge, SecurityControl, Vulnerability35from seeder import _ASSET_CATALOGUE, _CONTROL_CATALOGUE, _generate_vulnerabilities_for_asset, _RNG_SEED36import random37 38# --------------------------------------------------------------------------- #39# Enterprise Network Topology Seed Data40# --------------------------------------------------------------------------- #41# Tuples: (source_node, target_node, weight, protocol, is_segmented)42ENTERPRISE_NETWORK_EDGES: list[tuple[str, str, float, str, bool]] = [43 # Perimeter to DMZ Gateways44 ("INTERNET", "upi-gateway-prod-01", 1.2, "HTTPS", False),45 ("INTERNET", "upi-gateway-prod-02", 1.2, "HTTPS", False),46 ("INTERNET", "mobile-banking-gw-01", 1.1, "HTTPS", False),47 ("INTERNET", "staff-portal-01", 1.5, "HTTPS", False),48 ("INTERNET", "staff-portal-02", 1.5, "HTTPS", False),49 50 # DMZ to Corporate & HR Services51 ("staff-portal-01", "call-center-crm-01", 1.8, "HTTPS", False),52 ("staff-portal-02", "erp-hr-01", 2.0, "HTTPS", False),53 ("erp-hr-01", "erp-finance-01", 2.4, "INTERNAL_RPC", False),54 55 # DMZ to Core Application Layer56 ("upi-gateway-prod-01", "trading-platform-01", 2.2, "INTERNAL_RPC", False),57 ("upi-gateway-prod-02", "trading-platform-01", 2.2, "INTERNAL_RPC", False),58 ("mobile-banking-gw-01", "trading-platform-01", 2.0, "INTERNAL_RPC", False),59 60 # Application Layer to Core Banking Systems (Segmented Crown Jewels)61 ("trading-platform-01", "core-bank-switch-01", 2.5, "INTERNAL_RPC", True),62 ("trading-platform-01", "core-bank-switch-02", 2.5, "INTERNAL_RPC", True),63 ("trading-platform-01", "cust-db-primary", 2.0, "SQL_NET", True),64 ("core-bank-switch-01", "cust-db-primary", 1.5, "INTERNAL_RPC", True),65 ("core-bank-switch-02", "cust-db-primary", 1.5, "INTERNAL_RPC", True),66 67 # Database Replication & Internal Analytics68 ("cust-db-primary", "cust-db-replica", 1.2, "DB_REPLICATION", True),69 ("call-center-crm-01", "cust-db-replica", 3.0, "HTTPS", True),70 ("erp-finance-01", "cust-db-primary", 2.8, "SQL_NET", True),71 72 # Branch Operations & Office Controllers73 ("office-wifi-controller", "branch-pos-north-12", 1.0, "WIFI_INTERNAL", False),74 ("office-wifi-controller", "branch-pos-south-07", 1.0, "WIFI_INTERNAL", False),75 ("office-wifi-controller", "branch-pos-west-03", 1.0, "WIFI_INTERNAL", False),76 ("branch-pos-north-12", "call-center-crm-01", 2.5, "VPN", False),77 ("branch-pos-south-07", "call-center-crm-01", 2.5, "VPN", False),78 ("branch-pos-west-03", "call-center-crm-01", 2.5, "VPN", False),79 80 # Disaster Recovery Site Replication81 ("core-bank-switch-01", "dr-site-core-switch", 3.5, "ASYNC_MIRROR", True),82 ("core-bank-switch-02", "dr-site-core-switch", 3.5, "ASYNC_MIRROR", True),83 84 # Third-Party Vendor Integrations & AI Pipeline Topology85 ("INTERNET", "payment-aggregator-api", 1.2, "HTTPS", False),86 ("payment-aggregator-api", "upi-gateway-prod-01", 1.4, "INTERNAL_RPC", False),87 ("INTERNET", "hr-cloud-saas-gateway", 1.5, "HTTPS", False),88 ("hr-cloud-saas-gateway", "erp-hr-01", 1.8, "HTTPS", False),89 ("trading-platform-01", "fraud-detection-model-01", 1.6, "INTERNAL_RPC", False),90 ("fraud-detection-model-01", "cust-db-primary", 1.8, "SQL_NET", True),91 92 # Cross-Border & Foreign Cloud Integrations93 ("INTERNET", "global-crm-cloud", 1.5, "HTTPS", False),94 ("global-crm-cloud", "call-center-crm-01", 2.0, "HTTPS", False),95 ("INTERNET", "offshore-payment-analytics", 1.3, "HTTPS", False),96 ("offshore-payment-analytics", "upi-gateway-prod-01", 1.8, "INTERNAL_RPC", False),97 98 # Data Provenance & Unauthorized Shadow Leak Vectors (A -> B -> C)99 ("cust-db-primary", "analytics-integration-gw", 1.8, "INTERNAL_RPC", True, True, True, "Authorized Shared", "cust-db-primary:PII,cust-db-primary:FINANCIAL"),100 ("analytics-integration-gw", "third-party-marketing-sync", 0.9, "HTTPS_DATA_EXPORT", False, False, False, "Unauthorized Delegation", "cust-db-primary:PII,cust-db-primary:FINANCIAL"),101]102 103 104def test_database_connection() -> bool:105 """106 Validates connectivity to the configured database engine and prints diagnostics.107 """108 print(f"[*] Testing connection to: {engine.url.render_as_string(hide_password=True)} ...")109 try:110 with engine.connect() as conn:111 result = conn.execute(text("SELECT 1;")).scalar()112 try:113 version_info = conn.execute(text("SELECT version();")).scalar()114 except Exception:115 version_info = "Unknown (non-PostgreSQL dialect)"116 117 if result == 1:118 print("[+] Connection successful!")119 print(f"[+] Engine Dialect : {engine.dialect.name}")120 print(f"[+] Server Version : {version_info}")121 return True122 else:123 print("[-] Unexpected query result during health check.")124 return False125 except OperationalError as exc:126 print("[-] Database connection failed: OperationalError")127 print(f" Details: {exc.orig if hasattr(exc, 'orig') else exc}")128 print(" Ensure PostgreSQL is running and DATABASE_URL is properly configured.")129 return False130 except SQLAlchemyError as exc:131 print(f"[-] SQLAlchemy Error during connection test: {exc}")132 return False133 134 135def drop_and_recreate_tables() -> None:136 """Drops all existing tables in the schema and recreates them from ORM metadata."""137 print("[*] Dropping all existing database tables...")138 Base.metadata.drop_all(bind=engine)139 print("[+] Existing tables dropped.")140 141 print("[*] Creating all schema tables from declarative models...")142 Base.metadata.create_all(bind=engine)143 print("[+] All tables successfully created.")144 145 146def seed_enterprise_data(db: Session) -> dict[str, int]:147 """148 Populates initial enterprise inventory, vulnerabilities, controls, and network edges.149 """150 rng = random.Random(_RNG_SEED)151 summary: dict[str, int] = {152 "assets": 0,153 "vulnerabilities": 0,154 "controls": 0,155 "network_edges": 0,156 }157 158 print("[*] Seeding assets and vulnerabilities...")159 for entry in _ASSET_CATALOGUE:160 hostname = entry[0]161 asset_type = entry[1]162 tier = entry[2]163 business_unit = entry[3]164 rev_per_min = entry[4]165 pii_count = entry[5]166 fin_count = entry[6]167 rbi_reg = entry[7]168 sebi_reg = entry[8]169 hops = entry[9]170 171 if len(entry) == 15:172 data_residency_country = entry[10]173 cross_border_transfer_enabled = entry[11]174 destination_countries = entry[12]175 transfer_legal_mechanism = entry[13]176 is_rbi_localization_compliant = entry[14]177 178 classification_type = "Third-Party SaaS" if "SaaS" in asset_type else ("Core IT" if "Core" in asset_type else "Cloud Infrastructure")179 is_iso27001 = True180 is_iso42001 = False181 is_third_party = True if "Third-Party" in asset_type or "offshore" in hostname else False182 vendor_name = "Global CRM Services Ltd" if "crm" in hostname else ("Offshore Analytics Singapore Pte" if "offshore" in hostname else None)183 vendor_risk_tier = "Tier-2 High" if tier == "Medium" else ("Tier-1 Critical" if tier == "Critical" else "Tier-3 Standard")184 soc2_attestation = False185 elif len(entry) >= 17:186 classification_type = entry[10]187 is_iso27001 = entry[11]188 is_iso42001 = entry[12]189 is_third_party = entry[13]190 vendor_name = entry[14]191 vendor_risk_tier = entry[15]192 soc2_attestation = entry[16]193 194 data_residency_country = entry[17] if len(entry) > 17 else "IN"195 cross_border_transfer_enabled = entry[18] if len(entry) > 18 else False196 destination_countries = entry[19] if len(entry) > 19 else None197 transfer_legal_mechanism = entry[20] if len(entry) > 20 else "None"198 is_rbi_localization_compliant = entry[21] if len(entry) > 21 else True199 else:200 classification_type = "Core IT"201 is_iso27001 = True202 is_iso42001 = False203 is_third_party = False204 vendor_name = None205 vendor_risk_tier = None206 soc2_attestation = False207 data_residency_country = "IN"208 cross_border_transfer_enabled = False209 destination_countries = None210 transfer_legal_mechanism = "None"211 is_rbi_localization_compliant = True212 213 asset = Asset(214 hostname=hostname,215 asset_type=asset_type,216 tier=tier,217 business_unit=business_unit,218 revenue_per_minute=rev_per_min,219 pii_records_count=pii_count,220 financial_records_count=fin_count,221 is_rbi_regulated=rbi_reg,222 is_sebi_regulated=sebi_reg,223 network_hops_from_internet=hops,224 classification_type=classification_type,225 is_iso27001_regulated=is_iso27001,226 is_iso42001_regulated=is_iso42001,227 is_third_party=is_third_party,228 vendor_name=vendor_name,229 vendor_risk_tier=vendor_risk_tier,230 soc2_attestation=soc2_attestation,231 data_residency_country=data_residency_country,232 cross_border_transfer_enabled=cross_border_transfer_enabled,233 destination_countries=destination_countries,234 transfer_legal_mechanism=transfer_legal_mechanism,235 is_rbi_localization_compliant=is_rbi_localization_compliant,236 )237 for vuln_data in _generate_vulnerabilities_for_asset(rng, tier):238 asset.vulnerabilities.append(Vulnerability(**vuln_data))239 summary["vulnerabilities"] += 1240 241 db.add(asset)242 summary["assets"] += 1243 244 print("[*] Seeding security controls catalogue...")245 for code, name, target_tier, cost_inr, likelihood_reduction in _CONTROL_CATALOGUE:246 db.add(247 SecurityControl(248 code=code,249 name=name,250 target_tier=target_tier,251 cost_inr=cost_inr,252 likelihood_reduction=likelihood_reduction,253 is_active=False,254 )255 )256 summary["controls"] += 1257 258 print("[*] Seeding network attack graph edges...")259 for entry in ENTERPRISE_NETWORK_EDGES:260 src, tgt, weight, proto, is_seg = entry[:5]261 is_auth = entry[5] if len(entry) > 5 else True262 consent = entry[6] if len(entry) > 6 else True263 flow_type = entry[7] if len(entry) > 7 else "Direct API"264 data_tags = entry[8] if len(entry) > 8 else None265 edge = NetworkEdge(266 source_node=src,267 target_node=tgt,268 weight=weight,269 protocol=proto,270 is_segmented=is_seg,271 is_authorized_flow=is_auth,272 consent_recorded=consent,273 flow_type=flow_type,274 data_tags=data_tags,275 )276 db.add(edge)277 summary["network_edges"] += 1278 279 print("[*] Seeding data lineage & provenance flows...")280 node_a = db.query(Asset).filter_by(hostname="cust-db-primary").first()281 node_b = db.query(Asset).filter_by(hostname="analytics-integration-gw").first()282 node_c = db.query(Asset).filter_by(hostname="third-party-marketing-sync").first()283 if node_a and node_b and node_c:284 flow = DataLineageFlow(285 origin_asset_id=node_a.id,286 intermediary_asset_id=node_b.id,287 destination_asset_id=node_c.id,288 is_authorized=False,289 has_user_consent=False,290 records_exposed_pii=node_a.pii_records_count,291 records_exposed_financial=node_a.financial_records_count,292 detection_status="Unauthorized Sub-Processing",293 )294 db.add(flow)295 summary["data_lineage_flows"] = 1296 297 db.commit()298 print(f"[+] Seeding complete: {summary}")299 return summary300 301 302def verify_attack_graph(db: Session) -> None:303 """Verifies that the attack graph engine correctly parses the seeded topology."""304 print("\n[*] Verifying attack graph Dijkstra path calculations...")305 engine_inst = AttackGraphEngine(session=db)306 crown_jewels = engine_inst.shortest_path_to_crown_jewels(source="INTERNET")307 308 print(f"[+] Attack graph nodes count : {engine_inst.graph.number_of_nodes()}")309 print(f"[+] Attack graph edges count : {engine_inst.graph.number_of_edges()}")310 print(f"[+] Critical Crown Jewels reachable: {len(crown_jewels)}")311 312 for asset_id, data in crown_jewels.items():313 path_str = " -> ".join(data["path"]) if data["path"] else "UNREACHABLE"314 print(f" • [{data['tier']}] {data['hostname']}: hops={data['hops']}, cost={data['dijkstra_distance']}")315 print(f" Path: {path_str}")316 317 318def parse_arguments() -> argparse.Namespace:319 """Configures CLI flags for database initialization."""320 parser = argparse.ArgumentParser(321 description="Bootstrap and seed the PostgreSQL database for the Cyber Risk Quantification platform."322 )323 parser.add_argument(324 "-f", "--force",325 action="store_true",326 help="Execute drop and recreate without interactive confirmation prompt.",327 )328 parser.add_argument(329 "--test-only",330 action="store_true",331 help="Test database connection without altering schemas or data.",332 )333 parser.add_argument(334 "--skip-seed",335 action="store_true",336 help="Create tables without populating seed data.",337 )338 return parser.parse_args()339 340 341def main() -> None:342 """Primary entry point for the database initialization workflow."""343 args = parse_arguments()344 345 print("=" * 76)346 print(" AI-Powered Cyber Risk Quantification Platform - Database Initializer")347 print("=" * 76)348 349 # 1. Test database connection350 if not test_database_connection():351 sys.exit(1)352 353 if args.test_only:354 print("[+] Test-only flag specified. Exiting cleanly.")355 sys.exit(0)356 357 # 2. Confirmation prompt358 if not args.force:359 confirm = input("\n[?] WARNING: This will DROP and RECREATE all tables. Continue? (y/N): ").strip().lower()360 if confirm not in ("y", "yes"):361 print("[-] Operation aborted by user.")362 sys.exit(0)363 364 # 3. Drop and recreate tables365 try:366 drop_and_recreate_tables()367 except SQLAlchemyError as exc:368 print(f"[-] Error dropping or creating tables: {exc}")369 sys.exit(1)370 371 # 4. Seed initial enterprise records372 if not args.skip_seed:373 with SessionLocal() as db:374 try:375 seed_enterprise_data(db)376 verify_attack_graph(db)377 except SQLAlchemyError as exc:378 db.rollback()379 print(f"[-] Database error during seeding: {exc}")380 sys.exit(1)381 382 print("\n[+] Database initialization completed successfully.")383 384 385if __name__ == "__main__":386 main()387