Merge pull request #707 from cisagov/cd_elasticache_updated

Add Stats endpoints, settings and management commands

.pre-commit-config.yaml

@@ -107,6 +107,7 @@ repos:
       - id: mypy
         additional_dependencies:
           - types-requests
+          - redis==5.2.0
   - repo: https://github.com/asottile/pyupgrade
     rev: v3.10.1
     hooks:

backend/env.yml

@@ -9,6 +9,7 @@ staging:
   DB_NAME: ${ssm:/crossfeed/staging/DATABASE_NAME}
   DB_USERNAME: ${ssm:/crossfeed/staging/DATABASE_USER}
   DB_PASSWORD: ${ssm:/crossfeed/staging/DATABASE_PASSWORD}
+  DJANGO_SECRET: ${ssm:/crossfeed/staging/DJANGO_SECRECT}
   MDL_USERNAME: ${ssm:/crossfeed/staging/MDL_USERNAME}
   MDL_PASSWORD: ${ssm:/crossfeed/staging/MDL_PASSWORD}
   MDL_NAME: ${ssm:/crossfeed/staging/MDL_NAME}
@@ -44,6 +45,7 @@ staging:
   WORKER_USER_AGENT: ${ssm:/crossfeed/staging/WORKER_USER_AGENT}
   WORKER_SIGNATURE_PUBLIC_KEY: ${ssm:/crossfeed/staging/WORKER_SIGNATURE_PUBLIC_KEY}
   ELASTICSEARCH_ENDPOINT: ${ssm:/crossfeed/staging/ELASTICSEARCH_ENDPOINT}
+  ELASTICACHE_ENDPOINT: ${ssm:/crossfeed/staging/ELASTICACHE_ENDPOINT}
   REACT_APP_TERMS_VERSION: ${ssm:/crossfeed/staging/REACT_APP_TERMS_VERSION}
   REACT_APP_RANDOM_PASSWORD: ${ssm:/crossfeed/staging/REACT_APP_RANDOM_PASSWORD}
   MATOMO_URL: http://matomo.crossfeed.local
@@ -70,6 +72,7 @@ prod:
   DB_PASSWORD: ${ssm:/crossfeed/prod/DATABASE_PASSWORD}
   MDL_USERNAME: ${ssm:/crossfeed/prod/MDL_USERNAME}
   MDL_PASSWORD: ${ssm:/crossfeed/prod/MDL_PASSWORD}
+  DJANGO_SECRET: ${ssm:/crossfeed/prod/DJANGO_SECRECT}
   MDL_NAME: ${ssm:/crossfeed/prod/MDL_NAME}
   MI_ACCOUNT_NAME: ${ssm:/readysetcyber/prod/MI_ACCOUNT_NAME}
   MI_PASSWORD: ${ssm:/readysetcyber/prod/MI_ACCOUNT_PASSWORD}
@@ -95,6 +98,7 @@ prod:
   WORKER_USER_AGENT: ${ssm:/crossfeed/prod/WORKER_USER_AGENT}
   WORKER_SIGNATURE_PUBLIC_KEY: ${ssm:/crossfeed/prod/WORKER_SIGNATURE_PUBLIC_KEY}
   ELASTICSEARCH_ENDPOINT: ${ssm:/crossfeed/prod/ELASTICSEARCH_ENDPOINT}
+  ELASTICACHE_ENDPOINT: ${ssm:/crossfeed/prod/ELASTICACHE_ENDPOINT}
   REACT_APP_TERMS_VERSION: ${ssm:/crossfeed/prod/REACT_APP_TERMS_VERSION}
   REACT_APP_RANDOM_PASSWORD: ${ssm:/crossfeed/prod/REACT_APP_RANDOM_PASSWORD}
   MATOMO_URL: http://matomo.crossfeed.local

backend/requirements.txt

@@ -3,13 +3,16 @@ boto3
 cryptography==38.0.0
 django
 docker
-elasticsearch==7.9.0
+elasticsearch
 fastapi==0.111.0
 mangum==0.17.0
 minio
+mypy
 psycopg2-binary
 PyJWT
 pytest
 pytest-django
+redis==5.2.0
 requests==2.32.3
+types-redis
 uvicorn==0.30.1

backend/src/tasks/functions.yml

@@ -58,3 +58,125 @@ updateScanTaskStatus:
           detail:
             clusterArn:
               - ${file(env.yml):${self:provider.stage}-ecs-cluster, ''}
+
+populateServiceStatsElasticache:
+  handler: src/xfd_django/xfd_api/tasks/elasticache_tasks.populate_ServicesStatscache
+  runtime: python3.11
+  environment:
+    DJANGO_SETTINGS_MODULE: xfd_django.settings
+    REDIS_ENDPOINT: ${env:ELASTICACHE_ENDPOINT}
+  events:
+    - schedule:
+        rate: ron(0 0 * * ? *)  # This triggers the function every day it mightnight
+        enabled: true
+  package:
+    include:
+      - src/**
+      - requirements.txt
+    exclude:
+      - node_modules/**
+
+populatePortStatsElasticache:
+  handler: src/xfd_django/xfd_api/tasks/elasticache_tasks.populate_PortsStatscache
+  runtime: python3.11
+  environment:
+    DJANGO_SETTINGS_MODULE: xfd_django.settings
+    REDIS_ENDPOINT: ${env:ELASTICACHE_ENDPOINT}
+  events:
+    - schedule:
+        rate: ron(0 0 * * ? *)  # This triggers the function every day it mightnight
+        enabled: true
+  package:
+    include:
+      - src/**
+      - requirements.txt
+    exclude:
+      - node_modules/**
+
+populateNumVulnsStatsElasticache:
+  handler: src/xfd_django/xfd_api/tasks/elasticache_tasks.populate_NumVulnerabilitiesStatscache
+  runtime: python3.11
+  environment:
+    DJANGO_SETTINGS_MODULE: xfd_django.settings
+    REDIS_ENDPOINT: ${env:ELASTICACHE_ENDPOINT}
+  events:
+    - schedule:
+        rate: ron(0 0 * * ? *)  # This triggers the function every day it mightnight
+        enabled: true
+  package:
+    include:
+      - src/**
+      - requirements.txt
+    exclude:
+      - node_modules/**
+
+populateLatestVulnsStatsElasticache:
+  handler: src/xfd_django/xfd_api/tasks/elasticache_tasks.populate_LatestVulnerabilitiesCache
+  runtime: python3.11
+  environment:
+    DJANGO_SETTINGS_MODULE: xfd_django.settings
+    REDIS_ENDPOINT: ${env:ELASTICACHE_ENDPOINT}
+  events:
+    - schedule:
+        rate: ron(0 0 * * ? *)  # This triggers the function every day it mightnight
+        enabled: true
+  package:
+    include:
+      - src/**
+      - requirements.txt
+    exclude:
+      - node_modules/**
+
+populateMostCommonVulnerabilityElasticache:
+  handler: src/xfd_django/xfd_api/tasks/elasticache_tasks.populate_MostCommonVulnerabilitiesCache
+
+  runtime: python3.11
+  environment:
+    DJANGO_SETTINGS_MODULE: xfd_django.settings
+    REDIS_ENDPOINT: ${env:ELASTICACHE_ENDPOINT}
+  events:
+    - schedule:
+        rate: ron(0 0 * * ? *)  # This triggers the function every day it mightnight
+        enabled: true
+  package:
+    include:
+      - src/**
+      - requirements.txt
+    exclude:
+      - node_modules/**
+
+populateSeverityCountsCache:
+  handler: src/xfd_django/xfd_api/tasks/elasticache_tasks.populate_SeverityCountsCache
+
+  runtime: python3.11
+  environment:
+    DJANGO_SETTINGS_MODULE: xfd_django.settings
+    REDIS_ENDPOINT: ${env:ELASTICACHE_ENDPOINT}
+  events:
+    - schedule:
+        rate: ron(0 0 * * ? *)  # This triggers the function every day it mightnight
+        enabled: true
+  package:
+    include:
+      - src/**
+      - requirements.txt
+    exclude:
+      - node_modules/**
+
+populateByOrgCache:
+  handler: src/xfd_django/xfd_api/tasks/elasticache_tasks.populate_VulnerabilitiesByOrgCache
+
+  runtime: python3.11
+  environment:
+    DJANGO_SETTINGS_MODULE: xfd_django.settings
+    REDIS_ENDPOINT: ${env:ELASTICACHE_ENDPOINT}
+  events:
+    - schedule:
+        rate: ron(0 0 * * ? *)  # This triggers the function every day it mightnight
+        enabled: true
+  package:
+    include:
+      - src/**
+      - requirements.txt
+    exclude:
+      - node_modules/**

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

@@ -3,7 +3,7 @@
 import json
 
 # Third-Party Libraries
-from fastapi import HTTPException, status
+from fastapi import HTTPException, Security, status
 from fastapi.responses import JSONResponse
 from xfd_api.auth import get_jwt_from_code, process_user
 

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

@@ -13,7 +13,7 @@
 from django.http import Http404
 from fastapi import HTTPException
 
-from ..auth import get_org_memberships, is_global_view_admin
+from ..auth import get_org_memberships, get_user_organization_ids, is_global_view_admin
 from ..helpers.filter_helpers import filter_domains, sort_direction
 from ..models import Domain
 from ..schema_models.domain import DomainFilters, DomainSearch
@@ -82,3 +82,62 @@ def export_domains(domain_filters: DomainFilters):
         raise HTTPException(status_code=404, detail=str(e))
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+
+
+async def stats_total_domains(organization, tag, current_user):
+    """
+    Get total number of domains
+    Returns:
+        int: total number of domains
+    """
+    try:
+        # Base QuerySet
+        queryset = Domain.objects.all()
+
+        # Apply filtering logic at the endpoint
+        # Check if the user is a global admin
+        is_admin = is_global_view_admin(current_user)
+
+        # Get user's accessible organizations
+        if not is_admin:
+            user_org_ids = await get_user_organization_ids(current_user)
+            if not user_org_ids:
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="User does not belong to any organizations.",
+                )
+            queryset = queryset.filter(organizationId__id__in=user_org_ids)
+        else:
+            user_org_ids = None  # Admin has access to all organizations
+
+        # Apply organization filter
+        if organization:
+            if user_org_ids is not None and organization not in user_org_ids:
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="User does not have access to the specified organization.",
+                )
+            queryset = queryset.filter(organizationId__id=organization)
+
+        # Apply tag filter
+        if tag:
+            tag_org_ids = get_tag_organization_ids(tag)
+            if user_org_ids is not None:
+                accessible_org_ids = set(user_org_ids).intersection(tag_org_ids)
+                if not accessible_org_ids:
+                    raise HTTPException(
+                        status_code=status.HTTP_403_FORBIDDEN,
+                        detail="No accessible organizations found for the specified tag.",
+                    )
+                queryset = queryset.filter(organizationId__id__in=accessible_org_ids)
+            else:
+                queryset = queryset.filter(organizationId__id__in=tag_org_ids)
+
+        # Get total count
+        total_domains = await sync_to_async(queryset.count)()
+
+        # Return the count in the expected schema
+        return {"value": total_domains}
+
+    except HTTPException as http_exc:
+        raise http_exc

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

@@ -10,6 +10,7 @@
 
 from ..auth import (
     get_org_memberships,
+    get_user_organization_ids,
     is_global_view_admin,
     is_global_write_admin,
     is_org_admin,
@@ -19,7 +20,7 @@
 )
 from ..helpers.regionStateMap import REGION_STATE_MAP
 from ..models import Organization, OrganizationTag, Role, Scan, ScanTask, User
-from ..schema_models import organization as organization_schemas
+from ..schema_models import organization_schema
 
 
 def is_valid_uuid(val: str) -> bool:
@@ -330,7 +331,7 @@ def get_all_regions(current_user):
 
 
 def find_or_create_tags(
-    tags: List[organization_schemas.TagSchema],
+    tags: List[organization_schema.TagSchema],
 ) -> List[OrganizationTag]:
     """Find or create organization tags."""
     final_tags = []
@@ -984,3 +985,86 @@ def list_organizations_v2(state, regionId, current_user):
     except Exception as e:
         print(e)
         raise HTTPException(status_code=500, detail=str(e))
+
+
+async def stats_get_org_count_by_id(organization, tag, current_user, redis_client):
+    try:
+        # Retrieve data from Redis
+        json_data = await redis_client.get("vulnerabilities_by_org")
+
+        if json_data is None:
+            raise HTTPException(status_code=404, detail="Data not found in cache.")
+
+        vulnerabilities_data = json.loads(json_data)
+
+        # Get user's organization IDs
+        user_org_ids = await get_user_organization_ids(current_user)
+        if not user_org_ids:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="User does not belong to any organizations.",
+            )
+
+        # Check if user is a global admin
+        is_admin = is_global_view_admin(current_user)
+
+        # Determine accessible organizations
+        if is_admin:
+            accessible_org_ids = None
+        else:
+            accessible_org_ids = set(user_org_ids)
+
+        # Apply filters
+        if organization:
+            if (
+                accessible_org_ids is not None
+                and organization not in accessible_org_ids
+            ):
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="User does not have access to the specified organization.",
+                )
+            accessible_org_ids = {organization}
+        elif tag:
+            tag_org_ids = get_tag_organization_ids(tag)
+            if accessible_org_ids is not None:
+                accessible_org_ids = accessible_org_ids.intersection(tag_org_ids)
+            else:
+                accessible_org_ids = set(tag_org_ids)
+            if not accessible_org_ids:
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="No accessible organizations found for the specified tag.",
+                )
+
+        # Filter vulnerabilities
+        if accessible_org_ids is not None:
+            filtered_vulnerabilities = [
+                vuln
+                for vuln in vulnerabilities_data
+                if vuln["orgId"] in accessible_org_ids
+            ]
+        else:
+            filtered_vulnerabilities = vulnerabilities_data
+
+        # Aggregate counts by organization
+        org_counts = {}
+        for vuln in filtered_vulnerabilities:
+            org_id = vuln["orgId"]
+            org_name = vuln["orgName"]
+            if org_id not in org_counts:
+                org_counts[org_id] = {
+                    "id": org_name,
+                    "orgId": org_id,
+                    "value": 0,
+                    "label": org_name,
+                }
+            org_counts[org_id]["value"] += 1
+
+        # Convert to list and sort
+        results = sorted(org_counts.values(), key=lambda x: x["value"], reverse=True)
+
+        return results
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))