Merge pull request #1416 from cisagov/Domain/IP-Filter-Refactor-CRASM-2704

Domain/IP Filter Refactor (CRASM-3594)

Author: cduhn17

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

@@ -4,7 +4,7 @@
 import csv
 import io
 import logging
-from typing import Optional
+from typing import Any, Dict, Optional
 
 # Third-Party Libraries
 from django.core.paginator import Paginator
@@ -13,12 +13,15 @@
 from django.db.models.fields import GenericIPAddressField
 from django.db.models.functions import Cast
 from fastapi import HTTPException, status
-from xfd_mini_dl.models import Domain, DomainSearchView, Organization, Service
+from xfd_mini_dl.models import Domain, DomainSearchView, Organization, Service, UserType
 
+from ..api_methods.organization import escape_special_characters
+from ..api_methods.search import is_valid_org, is_valid_region
 from ..auth import get_org_memberships, is_global_view_admin
 from ..helpers.filter_helpers import apply_domain_filters
 from ..helpers.s3_client import S3Client
-from ..schema_models.domain import DomainSearch
+from ..schema_models.domain import DomainNameSearch, DomainSearch
+from ..tasks.es_client import ESClient
 
 LOGGER = logging.getLogger(__name__)
 
@@ -316,3 +319,99 @@ def export_domains(domain_search: DomainSearch, current_user):
         # Log the exception for debugging (optional)
         LOGGER.error("Error exporting domains: %s", e)
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
+
+def search_domains_name(search_body: DomainNameSearch, current_user):
+    """Handle the logic for searching organizations in Elasticsearch."""
+    try:
+        if search_body.regions is not None and len(search_body.regions) > 0:
+            # Validate regions
+            for region in search_body.regions:
+                if not is_valid_region(region, current_user):
+                    raise HTTPException(
+                        status_code=403,
+                        detail="Unauthorized",
+                    )
+        if search_body.organizations is not None and len(search_body.organizations) > 0:
+            # Validate organizations
+            for org in search_body.organizations:
+                if not is_valid_org(org, current_user):
+                    raise HTTPException(
+                        status_code=403,
+                        detail="Unauthorized",
+                    )
+        # Check if user is GlobalViewAdmin or has memberships
+        if not is_global_view_admin(current_user) and not get_org_memberships(
+            current_user
+        ):
+            return []
+
+        # Initialize Elasticsearch client
+        client = ESClient()
+
+        # Construct the Elasticsearch query
+
+        query_body: Dict[str, Any] = {
+            # "_source": ["id", "name"],
+            "query": {"bool": {"must": [], "filter": []}},
+        }
+
+        validated_search_field = (
+            search_body.search_field
+            if search_body.search_field in ["name", "ip"]
+            else "name"
+        )
+
+        # Use match_all if searchTerm is empty
+        if search_body.search_term.strip():
+            sanitized_search_term = escape_special_characters(search_body.search_term)
+            query_body["query"]["bool"]["must"].append(
+                {
+                    "query_string": {
+                        "query": "*{}*".format(sanitized_search_term),
+                        "fields": [validated_search_field],
+                        "fuzziness": "AUTO",
+                        "analyze_wildcard": True,
+                    }
+                }
+            )
+        else:
+            query_body["query"]["bool"]["must"].append({"match_all": {}})
+
+        # Apply region filters if provided
+        if search_body.regions:
+            query_body["query"]["bool"]["filter"].append(
+                {"terms": {"organization.region_id": search_body.regions}}
+            )
+        if search_body.organizations:
+            query_body["query"]["bool"]["filter"].append(
+                {"terms": {"organization.id.keyword": search_body.organizations}}
+            )
+
+        if current_user.user_type == UserType.STANDARD:
+            if search_body.regions == [] and search_body.organizations == []:
+                orgs = get_org_memberships(current_user)
+                if not orgs:
+                    return []
+                query_body["query"]["bool"]["filter"].append(
+                    {"terms": {"organization.id.keyword": orgs}}
+                )
+
+        if current_user.user_type == UserType.REGIONAL_ADMIN:
+            if search_body.regions == [] and search_body.organizations == []:
+                query_body["query"]["bool"]["filter"].append(
+                    {"terms": {"organization.region_id": [current_user.region_id]}}
+                )
+
+        # Log the query for debugging
+        LOGGER.debug("Query body: %s", query_body)
+
+        # Execute the search
+        search_results = client.search_domains(query_body)
+
+        return {"body": search_results}
+    except HTTPException as http_exc:
+        raise http_exc
+    except Exception as e:
+        LOGGER.exception("Error occurred while searching organizations: %s", e)
+        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)

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

@@ -5,6 +5,7 @@
 import socket
 
 # Third-Party Libraries
+from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
 import dns.resolver
 from xfd_api.helpers.asset_inserts import create_or_update_ip
@@ -16,6 +17,8 @@
 
 def get_matching_cidr(ip, org):
     """Return cidr that contains the ip owned by the org."""
+    if settings.IS_LOCAL:
+        return Cidr.objects.filter().first()
     try:
         # Use .get() to find a single CIDR network that contains the IP
         matching_cidr = Cidr.objects.get(
@@ -35,6 +38,12 @@ def get_matching_cidr(ip, org):
 def resolve_domain(domain, nameservers=None):
     """Identify ips linked to a given domain."""
     ip_addresses = set()
+
+    # If local add ip address based on string value instead of lookup
+    if settings.IS_LOCAL:
+        ip_addresses.add((domain.ip_address, "IPv4"))
+        return ip_addresses
+
     if not nameservers:
         nameservers = ["8.8.8.8"]
     # Create a resolver instance and optionally set a custom DNS server
@@ -46,7 +55,7 @@ def resolve_domain(domain, nameservers=None):
 
     try:
         # Resolve IPv4 addresses (A records)
-        ipv4_answers = dns.resolver.resolve(domain, "A")
+        ipv4_answers = dns.resolver.resolve(domain.sub_domain, "A")
         for rdata in ipv4_answers:
             ip_addresses.add((rdata.address, "IPv4"))
     except dns.resolver.NoAnswer:
@@ -56,7 +65,7 @@ def resolve_domain(domain, nameservers=None):
 
     try:
         # Resolve IPv6 addresses (AAAA records)
-        ipv6_answers = dns.resolver.resolve(domain, "AAAA")
+        ipv6_answers = dns.resolver.resolve(domain.sub_domain, "AAAA")
         for rdata in ipv6_answers:
             ip_addresses.add((rdata.address, "IPv6"))
     except dns.resolver.NoAnswer:
@@ -74,9 +83,10 @@ def get_ips_and_type_dns(subdomain, org):
     for ip_address, version in ip_set:
         cidr = get_matching_cidr(ip_address, org)
         if cidr:
-            LOGGER.warning(
-                "Found matching cidr for %s: %s", str(ip_address), cidr.network
-            )
+            if not settings.IS_LOCAL:
+                LOGGER.warning(
+                    "Found matching cidr for %s: %s", str(ip_address), cidr.network
+                )
             ip_info.append((ip_address, version, cidr))
     return ip_info
 
@@ -121,7 +131,7 @@ def get_ips_and_type_socket(subdomain, org):
 
 def link_ip_from_domain(sub, org):
     """Link IP from domain."""
-    ips = get_ips_and_type_dns(sub.sub_domain, org)
+    ips = get_ips_and_type_dns(sub, org)
 
     if not ips:
         return 0

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

@@ -2,7 +2,7 @@
 # Standard Python Libraries
 from datetime import datetime
 import logging
-from typing import Any, List, Optional
+from typing import Any, List, Literal, Optional
 from uuid import UUID
 
 # Third-Party Libraries
@@ -57,6 +57,15 @@ class Config:
         from_attributes = True
 
 
+class DomainNameSearch(BaseModel):
+    """DomainNameSearch Schema."""
+
+    regions: Optional[List[str]]
+    organizations: Optional[List[str]]
+    search_term: str
+    search_field: Literal["name", "ip"] = "name"
+
+
 class DomainSearch(BaseModel):
     """DomainSearch schema."""
 

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

@@ -21,8 +21,9 @@
 from django.db import IntegrityError, transaction
 from django.utils import timezone
 from faker import Faker
+from xfd_api.helpers.link_ips_from_subs import connect_ips_from_subs
 from xfd_api.helpers.regionStateMap import REGION_STATE_MAP, STATE_ABBR_MAP
-from xfd_api.models import Domain, Service, Vulnerability
+from xfd_api.models import Service, Vulnerability
 from xfd_api.schema_models.scan import SCAN_SCHEMA
 from xfd_api.tasks.refresh_material_views import handler as refresh_materialized_views
 from xfd_api.tasks.refresh_vs_summaries import handler as refresh_vs_summaries
@@ -33,6 +34,7 @@
     CidrOrgs,
     Cve,
     CveSsvc,
+    DataSource,
     Host,
     Ip,
     LatestPortScan,
@@ -42,6 +44,7 @@
     Role,
     Scan,
     ScanResult,
+    SubDomains,
     Ticket,
     TicketEvent,
     User,
@@ -751,6 +754,17 @@ def populate_sample_data():
         cidrs = generate_cidr_blocks()
         create_cidrs_for_org(org, cidrs)
 
+        data_source, _ = DataSource.objects.using("mini_data_lake").get_or_create(
+            name="findomain",
+            description="findomain enumerates domains into subs.",
+            last_run=datetime.now(),
+        )
+        for _ in range(NUM_SAMPLE_DOMAINS):
+            create_sample_domain(org, data_source)
+
+    # COnnect created sub domains with ips
+    connect_ips_from_subs(orgs)
+
     LOGGER.info("Populating vuln_scans, port_scans, tickets, and ticket_events...")
     for idx, org in enumerate(orgs, start=1):
         try:
@@ -1067,18 +1081,20 @@ def generate_random_name():
     return "{} {} {}".format(adjective.capitalize(), entity, noun.capitalize())
 
 
-def create_sample_domain(organization):
+def create_sample_domain(organization, data_source):
     """Create a sample domain linked to an organization."""
     domain_name = "{}-{}.crossfeed.local".format(
         random.choice(adjectives), random.choice(nouns)
     ).lower()
     ip = ".".join(map(str, (random.randint(0, 255) for _ in range(4))))
-    return Domain.objects.create(
-        name=domain_name,
-        ip=ip,
-        fromRootDomain="crossfeed.local",
-        subdomainSource="findomain",
+    return SubDomains.objects.create(
+        sub_domain=domain_name,
+        ip_address=ip,
+        from_root_domain="crossfeed.local",
+        subdomain_source="findomain",
         organization=organization,
+        data_source=data_source,
+        current=True,
     )
 
 

backend/src/xfd_django/xfd_api/tests/test_api_integrity.py

@@ -79,6 +79,7 @@ def test_endpoints_require_auth(method, route):
     ("DELETE", "/notifications/{notification_id}"),
     ("POST", "/v2/organizations/{organization_id}/users"),
     ("POST", "/search/organizations"),
+    ("POST", "/search/domains"),
     ("DELETE", "/saved-searches/{saved_search_id}"),
     ("POST", "/scheduler/invoke"),
     ("POST", "/scan-tasks/{scan_task_id}/kill"),

backend/src/xfd_django/xfd_api/tests/test_domain.py

@@ -3,6 +3,7 @@
 from datetime import datetime
 import logging
 import secrets
+from unittest.mock import patch
 
 # Third-Party Libraries
 from django.db import transaction
@@ -329,3 +330,80 @@ def test_search_domains_does_not_exist(user, domain, refresh_vuln_views):
     assert response.status_code == 200
     data = response.json()
     assert len(data["result"]) == 0, "No result found for the given organization name"
+
+
+@pytest.mark.django_db(transaction=True, databases=["default", "mini_data_lake"])
+@patch("xfd_api.tasks.es_client.ESClient.search_domains")
+def test_domains_search_autofill_endpoint_auth_user_200(mock_search):
+    """Test domain search autocomplete endpoint."""
+    user = User.objects.create(
+        first_name="",
+        last_name="",
+        email="{}@example.com".format(secrets.token_hex(4)),
+        user_type=UserType.STANDARD,
+        created_at=datetime.now(),
+        updated_at=datetime.now(),
+    )
+    response = client.post(
+        "/search/domains",
+        json={
+            "search_term": "127",
+            "search_field": "name",
+            "regions": [],
+            "organizations": [],
+        },
+        headers={"Authorization": "Bearer " + create_jwt_token(user)},
+    )
+    assert response.status_code == 200
+
+
+@pytest.mark.django_db(transaction=True, databases=["default", "mini_data_lake"])
+@patch("xfd_api.tasks.es_client.ESClient.search_domains")
+def test_domains_search_autofill_endpoint_regional_auth(mock_search):
+    """Test domain search autocomplete endpoint."""
+    user = User.objects.create(
+        first_name="",
+        last_name="",
+        email="{}@example.com".format(secrets.token_hex(4)),
+        user_type=UserType.REGIONAL_ADMIN,
+        created_at=datetime.now(),
+        updated_at=datetime.now(),
+        region_id="8",
+    )
+    response = client.post(
+        "/search/domains",
+        json={
+            "search_term": "127",
+            "search_field": "name",
+            "regions": ["3"],
+            "organizations": [],
+        },
+        headers={"Authorization": "Bearer " + create_jwt_token(user)},
+    )
+    assert response.status_code == 200
+
+
+@pytest.mark.django_db(transaction=True, databases=["default", "mini_data_lake"])
+@patch("xfd_api.tasks.es_client.ESClient.search_domains")
+def test_domains_search_autofill_endpoint_global_auth(mock_search):
+    """Test domain search autocomplete endpoint."""
+    user = User.objects.create(
+        first_name="",
+        last_name="",
+        email="{}@example.com".format(secrets.token_hex(4)),
+        user_type=UserType.GLOBAL_ADMIN,
+        created_at=datetime.now(),
+        updated_at=datetime.now(),
+        region_id="8",
+    )
+    response = client.post(
+        "/search/domains",
+        json={
+            "search_term": "127",
+            "search_field": "name",
+            "regions": ["3"],
+            "organizations": [],
+        },
+        headers={"Authorization": "Bearer " + create_jwt_token(user)},
+    )
+    assert response.status_code == 200