Merge pull request #922 from cisagov/DJ_optimize-asm-sync_WIP

optimize asm_sync to enumerate live ips

Author: rapidray12

backend/src/xfd_django/xfd_api/api_methods/sync.py

@@ -48,14 +48,14 @@ async def sync_post(sync_body, request: Request, current_user):
 
     except HTTPException as http_exc:
         raise http_exc
-    except Exception as e:
-        print(e)
-        raise HTTPException(status_code=500, detail=str(e))
     except SyncError as sync_exc:
         raise HTTPException(
             status_code=500,
             detail=f"SyncError: {sync_exc.message} - {sync_exc.error_message}",
         )
+    except Exception as e:
+        print(e)
+        raise HTTPException(status_code=500, detail=str(e))
 
 
 def process_request(headers, sync_body):

backend/src/xfd_django/xfd_api/helpers/asset_inserts.py

@@ -39,6 +39,9 @@ def create_or_update_ip(create_defaults, update_dict, linked_sub=None):
     if not created:
         for key, value in update_dict.items():
             if key == "origin_cidr":
+                if ip_object.origin_cidr is None:
+                    ip_object.origin_cidr = value
+                    continue
                 if value.id == ip_object.origin_cidr.id:
                     continue
                 if ip_object.origin_cidr.retired is True:

backend/src/xfd_django/xfd_api/helpers/link_subs_from_ips.py

@@ -35,51 +35,76 @@
 )
 
 
-def process_ips(thread_id, org, cidr, ip_gen):
-    """Process ips through WhoisXML and save them to DB."""
+def process_ips(thread_id, org, cidr, ip_list):
+    """Process IPs through WhoisXML and save them to DB."""
     count = 0
     failed_ips = []
     chunk_start = time.time()
-    while True:
+    for ip in ip_list:
+        count += 1
         try:
-            # Get the next IP from the generator or break if exhausted
-            ip = str(next(ip_gen))
-            count += 1
-            try:
-                domain_list, failed_ips = search_whois_for_domains(ip, failed_ips)
-            except Exception as e:
-                LOGGER.error("Thread %d: Error identifying domains: %s", thread_id, e)
-
-                failed_ips.append(ip)
-                continue
-            if domain_list:
-                LOGGER.info("Found %d domains associated with %s", len(domain_list), ip)
-                save_and_link_ip_and_subdomain(ip, cidr, org, domain_list)
-
-        except StopIteration:
-            # Stop when the generator is exhausted
-            LOGGER.info(
-                "Thread %d has completed. Processed %d ips in %s seconds.",
-                thread_id,
-                count,
-                round(time.time() - chunk_start, 2),
+            domain_list, failed_ips = search_whois_for_domains(ip, failed_ips)
+        except Exception as e:
+            LOGGER.error(
+                "Thread %d: Error identifying domains for %s: %s", thread_id, ip, e
             )
-            if len(failed_ips) > 0:
-                LOGGER.warning("%d IPs failed to process", len(failed_ips))
-            break
+            failed_ips.append(ip)
+            continue
+
+        if domain_list:
+            LOGGER.info("Found %d domains associated with %s", len(domain_list), ip)
+            save_and_link_ip_and_subdomain(ip, cidr, org, domain_list)
+
+    LOGGER.info(
+        "Thread %d completed. Processed %d IPs in %.2f seconds.",
+        thread_id,
+        count,
+        time.time() - chunk_start,
+    )
+
+    if failed_ips:
+        LOGGER.warning(
+            "%d IPs failed to process in thread %d", len(failed_ips), thread_id
+        )
+
+
+def split_into_balanced_chunks(items, num_chunks):
+    """Split `items` into `num_chunks` parts, distributing remainder fairly."""
+    n = len(items)
+    base_chunk_size = n // num_chunks
+    remainder = n % num_chunks
+
+    chunks = []
+    start = 0
+    for i in range(num_chunks):
+        # Give one extra item to the first `remainder` chunks
+        end = start + base_chunk_size + (1 if i < remainder else 0)
+        chunks.append(items[start:end])
+        start = end
+    return chunks
 
 
 def process_cidr(cidr, org):
-    """Process a given cidr."""
-    ip_gen = generate_ips(cidr.network)
+    """Process a given CIDR using stored live IPs."""
+    if not cidr.live_ips:
+        LOGGER.warning("No live IPs for CIDR: %s", cidr.network)
+        return
+
+    ip_list = list(cidr.live_ips)
+    if not ip_list:
+        return
+
+    chunks = split_into_balanced_chunks(ip_list, THREAD_COUNT)
 
     threads = []
-    for i in range(THREAD_COUNT):
-        thread = threading.Thread(target=process_ips, args=(i + 1, org, cidr, ip_gen))
-        threads.append(thread)
-        thread.start()
+    for i, chunk in enumerate(chunks):
+        if chunk:  # Only create a thread if the chunk has work
+            thread = threading.Thread(
+                target=process_ips, args=(i + 1, org, cidr, chunk)
+            )
+            threads.append(thread)
+            thread.start()
 
-    # Wait for all threads to complete
     for thread in threads:
         thread.join()
 

backend/src/xfd_django/xfd_api/management/commands/syncmdl.py

@@ -92,6 +92,15 @@ def handle(self, *args, **options):
             except Exception as e:
                 self.stdout.write("Granting privileges failed: {}".format(e))
 
+        # 👉 Step 1.5: Enable btree_gist extension
+        self.stdout.write("Enabling btree_gist extension for GiST indexing...")
+        try:
+            with connection.cursor() as cursor:
+                cursor.execute("CREATE EXTENSION IF NOT EXISTS btree_gist;")
+                self.stdout.write("btree_gist extension enabled.")
+        except Exception as e:
+            self.stdout.write(f"Failed to enable btree_gist extension: {e}")
+
         # Step 2: Synchronize or Reset the Database
         self.stdout.write("Synchronizing the MDL database schema...")
         if dangerouslyforce:

backend/src/xfd_django/xfd_api/schema_models/scan.py

@@ -147,7 +147,7 @@ class GenericMessageResponseModel(BaseModel):
         cpu="1024",
         memory="8192",
         description="Enumerate and sync org assets.",
-        max_concurrent_tasks=1,
+        max_concurrent_tasks=3,
     ),
     "censys": ScanSchema(
         type="fargate",

backend/src/xfd_django/xfd_api/tasks/syncdb_helpers.py

@@ -28,6 +28,9 @@
 from xfd_api.helpers.regionStateMap import REGION_STATE_MAP
 from xfd_api.models import Domain, Service, Vulnerability
 from xfd_api.tasks.es_client import ESClient
+from xfd_api.utils.scan_utils.vuln_scanning_sync_utils import (  # fill_cidr_live_ips,
+    fill_cidr_live_ips_bulk_update,
+)
 from xfd_mini_dl.models import (
     ApiKey,
     Cidr,
@@ -668,7 +671,7 @@ def create_cidrs_for_org(org, cidr_list, data_source=None, ips_per_cidr=4):
                     last_seen_timestamp=timezone.now(),
                     last_reverse_lookup=timezone.now(),
                     has_shodan_results=random.choice([True, False]),
-                    current=random.choice([True, False]),
+                    current=random.choice([True, True, True, False]),
                     conflict_alerts=[],
                     synced_at=timezone.now(),
                 )
@@ -741,6 +744,9 @@ def populate_sample_data():
         )
         sys.stdout.flush()
 
+    # fill_cidr_live_ips()
+    fill_cidr_live_ips_bulk_update()
+
     print("\n✅ Done populating all data.")
 
 
@@ -959,6 +965,22 @@ def synchronize(target_app_label=None, using=None):
         process_m2m_tables(schema_editor, ordered_models, database)
 
         if target_app_label == "xfd_mini_dl":
+            print("Ensuring GiST index exists on ip.ip...")
+            with connections[database].cursor() as cursor:
+                cursor.execute(
+                    """
+                    DO $$
+                    BEGIN
+                        IF NOT EXISTS (
+                            SELECT 1 FROM pg_indexes
+                            WHERE tablename = 'ip' AND indexname = 'ip_ip_gist_idx'
+                        ) THEN
+                            EXECUTE 'CREATE INDEX ip_ip_gist_idx ON ip USING gist (ip inet_ops)';
+                        END IF;
+                    END
+                    $$;
+                """
+                )
             create_domain_view(database)
             create_service_view(database)
             create_vuln_normal_views(database)
@@ -1039,6 +1061,7 @@ def process_model(
             else:
                 print("Creating table for model: {}".format(model.__name__))
                 schema_editor.create_model(model)
+
         except Exception as e:
             print("Error processing model {}: {}".format(model.__name__, e))
 

backend/src/xfd_django/xfd_api/tasks/vulnScanningSync.py

@@ -39,9 +39,10 @@
     ScanExecutionError,
     SyncError,
 )
-from xfd_api.utils.scan_utils.vuln_scanning_sync_utils import (
+from xfd_api.utils.scan_utils.vuln_scanning_sync_utils import (  # fill_cidr_live_ips,
     enforce_latest_flag_port_scan,
     fetch_orgs_and_relations,
+    fill_cidr_live_ips_bulk_update,
     get_latest_os_type,
     load_test_data,
     save_cve_to_datalake,
@@ -138,7 +139,7 @@ def fetch_in_chunks(base_query: str, chunk_size: int = 5000):
         offset += chunk_size
 
 
-def main():
+def main():  # pylint: disable=R0915
     """Execute the vulnerability scanning synchronization task."""
     LOGGER.info("Started VulnScanningSync scan...")
 
@@ -150,9 +151,6 @@ def main():
     org_id_dict = process_orgs(request_list)
     LOGGER.info("Completed saving organizations to the LZ MDL.")
 
-    # Process Organizations & Relations
-    send_organizations_to_dmz()
-
     # Process Vulnerability Scans
     LOGGER.info("Started processing vulnerability scans...")
     vuln_scans = fetch_from_redshift(
@@ -199,6 +197,12 @@ def main():
         create_port_scan_service_summaries()
         LOGGER.info("Finished processing port scans")
 
+    # fill_cidr_live_ips()
+    fill_cidr_live_ips_bulk_update()
+
+    # Process Organizations & Relations
+    send_organizations_to_dmz()
+
     # Process Tickets (Chunked)
     LOGGER.info("Started processing tickets...")
     base_query = (
@@ -343,7 +347,10 @@ def send_csv_to_sync(csv_data, bounds):
 
     try:
         response = requests.post(
-            os.getenv("DMZ_SYNC_ENDPOINT"), json=body, headers=headers, timeout=60
+            os.getenv("DMZ_SYNC_ENDPOINT") + "/sync",
+            json=body,
+            headers=headers,
+            timeout=60,
         )
         response.raise_for_status()
         LOGGER.info("Successfully sent chunk to sync API")

backend/src/xfd_django/xfd_api/utils/scan_utils/vuln_scanning_sync_utils.py

@@ -10,6 +10,7 @@
 import json
 import logging
 import os
+import time
 from typing import Dict
 from uuid import uuid1
 
@@ -18,6 +19,7 @@
 from django.db import connections, models, transaction
 from django.db.models import Exists, OuterRef, Prefetch
 from django.db.utils import IntegrityError
+from django.utils import timezone
 from xfd_mini_dl.models import (
     Cidr,
     CidrOrgs,
@@ -584,6 +586,7 @@ def save_cidr_to_mdl(cidr_dict: dict, org: Organization, db_name="mini_data_lake
                 cidr_obj.start_ip = cidr_dict["start_ip"]
                 cidr_obj.end_ip = cidr_dict["end_ip"]
                 cidr_obj.retired = False
+                cidr_obj.live_ips = cidr_dict["live_ips"]
                 cidr_obj.save(using=db_name)  # Save updates
 
             else:
@@ -592,6 +595,7 @@ def save_cidr_to_mdl(cidr_dict: dict, org: Organization, db_name="mini_data_lake
                     network=cidr_dict["network"],
                     start_ip=cidr_dict["start_ip"],
                     end_ip=cidr_dict["end_ip"],
+                    live_ips=cidr_dict["live_ips"],
                     retired=False,
                 )
             # cidr_obj.organizations.add(org, through_defaults={})
@@ -690,3 +694,86 @@ def map_severity(severity):
     if severity < 9:
         return "High"
     return "Critical"
+
+
+def fill_cidr_live_ips():
+    """Update live_ips field for all current CIDRs based on recent open PortScans."""
+    start_time = time.time()
+
+    # Define the 90-day threshold
+    time_threshold = timezone.now() - datetime.timedelta(days=90)
+
+    # Get all Cidrs with at least one related CidrOrgs marked as current
+    current_cidrs = Cidr.objects.filter(cidrorgs__current=True).distinct()
+
+    for cidr in current_cidrs:
+        if not cidr.network:
+            continue
+
+        scans = (
+            PortScan.objects.filter(
+                state="open",
+                time_scanned__gte=time_threshold,
+                ip__ip__net_contained=cidr.network,
+            )
+            .values_list("ip__ip", flat=True)
+            .distinct()
+        )
+
+        # If live_ips is empty or not set, initialize it as an empty set
+        current_live_ips = set(cidr.live_ips or [])
+
+        # Add the new IPs from the scans to the existing set (no duplicates)
+        current_live_ips.update(scans)
+
+        # Convert all IP objects to strings for JSON serialization
+        cidr.live_ips = [str(ip.ip) for ip in current_live_ips]
+        cidr.save()
+
+    duration = time.time() - start_time
+    LOGGER.info("fill_cidr_live_ips completed in %.2f seconds", duration)
+
+
+def fill_cidr_live_ips_bulk_update():
+    """Fill live_ips field in the cidr table based on recent port scans."""
+    start_time = time.time()
+
+    with transaction.atomic(using="mini_data_lake"):
+        with connections["mini_data_lake"].cursor() as cursor:
+            cursor.execute(
+                """
+                WITH new_ips AS (
+                    SELECT
+                        cidr.id AS cidr_id,
+                        array_agg(DISTINCT ip.ip) AS new_ip_list
+                    FROM cidr
+                    JOIN cidr_orgs ON cidr_orgs.cidr_id = cidr.id
+                    JOIN port_scan ON port_scan.state = 'open'
+                        AND port_scan.time_scanned >= NOW() - INTERVAL '90 days'
+                    JOIN ip ON port_scan.ip_id = ip.id
+                    WHERE cidr_orgs.current = TRUE
+                      AND cidr.network IS NOT NULL
+                      AND ip.ip << cidr.network
+                    GROUP BY cidr.id
+                ),
+                merged_ips AS (
+                    SELECT
+                        cidr.id,
+                        ARRAY(
+                            SELECT DISTINCT ip_address::inet
+                            FROM jsonb_array_elements_text(
+                                COALESCE(cidr.live_ips, '[]'::jsonb) || to_jsonb(new_ips.new_ip_list)
+                            ) AS ip_address
+                        ) AS updated_ips
+                    FROM cidr
+                    JOIN new_ips ON cidr.id = new_ips.cidr_id
+                )
+                UPDATE cidr
+                SET live_ips = to_jsonb(merged_ips.updated_ips)
+                FROM merged_ips
+                WHERE cidr.id = merged_ips.id;
+                """
+            )
+
+    duration = time.time() - start_time
+    LOGGER.info("fill_cidr_live_ips_bulk_update completed in %.2f seconds", duration)