stage 1 of doing new cidrs/ site setup

This commit is contained in:
2025-11-19 13:39:27 -06:00
parent 4a4c33a10b
commit 034f146fa1
16 changed files with 3998 additions and 609 deletions

View File

@@ -1,8 +1,8 @@
"""
Config Service - Business logic for config file management
Config Service - Business logic for config management
This service handles all operations related to scan configuration files,
including creation, validation, listing, and deletion.
This service handles all operations related to scan configurations,
both database-stored (primary) and file-based (deprecated).
"""
import os
@@ -13,26 +13,343 @@ from typing import Dict, List, Tuple, Any, Optional
from datetime import datetime
from pathlib import Path
from werkzeug.utils import secure_filename
from sqlalchemy.orm import Session
class ConfigService:
"""Business logic for config management"""
def __init__(self, configs_dir: str = '/app/configs'):
def __init__(self, db_session: Session = None, configs_dir: str = '/app/configs'):
"""
Initialize the config service.
Args:
configs_dir: Directory where config files are stored
db_session: SQLAlchemy database session (for database operations)
configs_dir: Directory where legacy config files are stored
"""
self.db = db_session
self.configs_dir = configs_dir
# Ensure configs directory exists
# Ensure configs directory exists (for legacy YAML configs)
os.makedirs(self.configs_dir, exist_ok=True)
def list_configs(self) -> List[Dict[str, Any]]:
# ============================================================================
# Database-based Config Operations (Primary)
# ============================================================================
def create_config(self, title: str, description: Optional[str], site_ids: List[int]) -> Dict[str, Any]:
"""
List all config files with metadata.
Create a new scan configuration in the database.
Args:
title: Configuration title
description: Optional configuration description
site_ids: List of site IDs to include in this config
Returns:
Created config as dictionary:
{
"id": 1,
"title": "Production Scan",
"description": "...",
"site_count": 3,
"sites": [...],
"created_at": "2025-11-19T10:30:00Z",
"updated_at": "2025-11-19T10:30:00Z"
}
Raises:
ValueError: If validation fails or sites don't exist
"""
if not title or not title.strip():
raise ValueError("Title is required")
if not site_ids or len(site_ids) == 0:
raise ValueError("At least one site must be selected")
# Import models here to avoid circular imports
from web.models import ScanConfig, ScanConfigSite, Site
# Verify all sites exist
existing_sites = self.db.query(Site).filter(Site.id.in_(site_ids)).all()
if len(existing_sites) != len(site_ids):
found_ids = {s.id for s in existing_sites}
missing_ids = set(site_ids) - found_ids
raise ValueError(f"Sites not found: {missing_ids}")
# Create config
config = ScanConfig(
title=title.strip(),
description=description.strip() if description else None,
created_at=datetime.utcnow(),
updated_at=datetime.utcnow()
)
self.db.add(config)
self.db.flush() # Get the config ID
# Create associations
for site_id in site_ids:
assoc = ScanConfigSite(
config_id=config.id,
site_id=site_id,
created_at=datetime.utcnow()
)
self.db.add(assoc)
self.db.commit()
return self.get_config_by_id(config.id)
def get_config_by_id(self, config_id: int) -> Dict[str, Any]:
"""
Get a scan configuration by ID.
Args:
config_id: Configuration ID
Returns:
Config as dictionary with sites
Raises:
ValueError: If config not found
"""
from web.models import ScanConfig
config = self.db.query(ScanConfig).filter_by(id=config_id).first()
if not config:
raise ValueError(f"Config with ID {config_id} not found")
# Get associated sites
sites = []
for assoc in config.site_associations:
site = assoc.site
sites.append({
'id': site.id,
'name': site.name,
'description': site.description,
'cidr_count': len(site.cidrs)
})
return {
'id': config.id,
'title': config.title,
'description': config.description,
'site_count': len(sites),
'sites': sites,
'created_at': config.created_at.isoformat() + 'Z' if config.created_at else None,
'updated_at': config.updated_at.isoformat() + 'Z' if config.updated_at else None
}
def list_configs_db(self) -> List[Dict[str, Any]]:
"""
List all scan configurations from database.
Returns:
List of config dictionaries with metadata
"""
from web.models import ScanConfig
configs = self.db.query(ScanConfig).order_by(ScanConfig.updated_at.desc()).all()
result = []
for config in configs:
sites = []
for assoc in config.site_associations:
site = assoc.site
sites.append({
'id': site.id,
'name': site.name
})
result.append({
'id': config.id,
'title': config.title,
'description': config.description,
'site_count': len(sites),
'sites': sites,
'created_at': config.created_at.isoformat() + 'Z' if config.created_at else None,
'updated_at': config.updated_at.isoformat() + 'Z' if config.updated_at else None
})
return result
def update_config(self, config_id: int, title: Optional[str], description: Optional[str], site_ids: Optional[List[int]]) -> Dict[str, Any]:
"""
Update a scan configuration.
Args:
config_id: Configuration ID to update
title: New title (optional)
description: New description (optional)
site_ids: New list of site IDs (optional, replaces existing)
Returns:
Updated config dictionary
Raises:
ValueError: If config not found or validation fails
"""
from web.models import ScanConfig, ScanConfigSite, Site
config = self.db.query(ScanConfig).filter_by(id=config_id).first()
if not config:
raise ValueError(f"Config with ID {config_id} not found")
# Update fields if provided
if title is not None:
if not title.strip():
raise ValueError("Title cannot be empty")
config.title = title.strip()
if description is not None:
config.description = description.strip() if description.strip() else None
# Update sites if provided
if site_ids is not None:
if len(site_ids) == 0:
raise ValueError("At least one site must be selected")
# Verify all sites exist
existing_sites = self.db.query(Site).filter(Site.id.in_(site_ids)).all()
if len(existing_sites) != len(site_ids):
found_ids = {s.id for s in existing_sites}
missing_ids = set(site_ids) - found_ids
raise ValueError(f"Sites not found: {missing_ids}")
# Remove existing associations
self.db.query(ScanConfigSite).filter_by(config_id=config_id).delete()
# Create new associations
for site_id in site_ids:
assoc = ScanConfigSite(
config_id=config_id,
site_id=site_id,
created_at=datetime.utcnow()
)
self.db.add(assoc)
config.updated_at = datetime.utcnow()
self.db.commit()
return self.get_config_by_id(config_id)
def delete_config(self, config_id: int) -> None:
"""
Delete a scan configuration from database.
This will cascade delete associated ScanConfigSite records.
Schedules and scans referencing this config will have their
config_id set to NULL.
Args:
config_id: Configuration ID to delete
Raises:
ValueError: If config not found
"""
from web.models import ScanConfig
config = self.db.query(ScanConfig).filter_by(id=config_id).first()
if not config:
raise ValueError(f"Config with ID {config_id} not found")
self.db.delete(config)
self.db.commit()
def add_site_to_config(self, config_id: int, site_id: int) -> Dict[str, Any]:
"""
Add a site to an existing config.
Args:
config_id: Configuration ID
site_id: Site ID to add
Returns:
Updated config dictionary
Raises:
ValueError: If config or site not found, or association already exists
"""
from web.models import ScanConfig, Site, ScanConfigSite
config = self.db.query(ScanConfig).filter_by(id=config_id).first()
if not config:
raise ValueError(f"Config with ID {config_id} not found")
site = self.db.query(Site).filter_by(id=site_id).first()
if not site:
raise ValueError(f"Site with ID {site_id} not found")
# Check if association already exists
existing = self.db.query(ScanConfigSite).filter_by(
config_id=config_id, site_id=site_id
).first()
if existing:
raise ValueError(f"Site '{site.name}' is already in this config")
# Create association
assoc = ScanConfigSite(
config_id=config_id,
site_id=site_id,
created_at=datetime.utcnow()
)
self.db.add(assoc)
config.updated_at = datetime.utcnow()
self.db.commit()
return self.get_config_by_id(config_id)
def remove_site_from_config(self, config_id: int, site_id: int) -> Dict[str, Any]:
"""
Remove a site from a config.
Args:
config_id: Configuration ID
site_id: Site ID to remove
Returns:
Updated config dictionary
Raises:
ValueError: If config not found, or removing would leave config empty
"""
from web.models import ScanConfig, ScanConfigSite
config = self.db.query(ScanConfig).filter_by(id=config_id).first()
if not config:
raise ValueError(f"Config with ID {config_id} not found")
# Check if this would leave the config empty
current_site_count = len(config.site_associations)
if current_site_count <= 1:
raise ValueError("Cannot remove last site from config. Delete the config instead.")
# Remove association
deleted = self.db.query(ScanConfigSite).filter_by(
config_id=config_id, site_id=site_id
).delete()
if deleted == 0:
raise ValueError(f"Site with ID {site_id} is not in this config")
config.updated_at = datetime.utcnow()
self.db.commit()
return self.get_config_by_id(config_id)
# ============================================================================
# Legacy YAML File Operations (Deprecated)
# ============================================================================
def list_configs_file(self) -> List[Dict[str, Any]]:
"""
[DEPRECATED] List all config files with metadata.
Returns:
List of config metadata dictionaries:
@@ -175,6 +492,9 @@ class ConfigService:
if not is_valid:
raise ValueError(f"Invalid config structure: {error_msg}")
# Create inline sites in database (if any)
self.create_inline_sites(parsed)
# Write file
with open(filepath, 'w') as f:
f.write(content)
@@ -266,9 +586,9 @@ class ConfigService:
return filename, yaml_content
def update_config(self, filename: str, yaml_content: str) -> None:
def update_config_file(self, filename: str, yaml_content: str) -> None:
"""
Update existing config file with new YAML content.
[DEPRECATED] Update existing config file with new YAML content.
Args:
filename: Config filename to update
@@ -299,9 +619,9 @@ class ConfigService:
with open(filepath, 'w') as f:
f.write(yaml_content)
def delete_config(self, filename: str) -> None:
def delete_config_file(self, filename: str) -> None:
"""
Delete config file and cascade delete any associated schedules.
[DEPRECATED] Delete config file and cascade delete any associated schedules.
When a config is deleted, all schedules using that config (both enabled
and disabled) are automatically deleted as well, since they would be
@@ -371,12 +691,15 @@ class ConfigService:
# Delete file
os.remove(filepath)
def validate_config_content(self, content: Dict) -> Tuple[bool, str]:
def validate_config_content(self, content: Dict, check_site_refs: bool = True) -> Tuple[bool, str]:
"""
Validate parsed YAML config structure.
Supports both legacy format (inline IPs) and new format (site references or CIDRs).
Args:
content: Parsed YAML config as dict
check_site_refs: If True, validates that referenced sites exist in database
Returns:
Tuple of (is_valid, error_message)
@@ -408,11 +731,65 @@ class ConfigService:
if not isinstance(site, dict):
return False, f"Site {i+1} must be a dictionary/object"
# Check if this is a site reference (new format)
if 'site_ref' in site:
# Site reference format
site_ref = site.get('site_ref')
if not isinstance(site_ref, str) or not site_ref.strip():
return False, f"Site {i+1} field 'site_ref' must be a non-empty string"
# Validate site reference exists (if check enabled)
if check_site_refs:
try:
from web.services.site_service import SiteService
from flask import current_app
site_service = SiteService(current_app.db_session)
referenced_site = site_service.get_site_by_name(site_ref)
if not referenced_site:
return False, f"Site {i+1}: Referenced site '{site_ref}' does not exist"
except Exception as e:
# If we can't check (e.g., outside app context), skip validation
pass
continue # Site reference is valid
# Check if this is inline site creation with CIDRs (new format)
if 'cidrs' in site:
# Inline site creation with CIDR format
if 'name' not in site:
return False, f"Site {i+1} with inline CIDRs missing required field: 'name'"
cidrs = site.get('cidrs')
if not isinstance(cidrs, list):
return False, f"Site {i+1} field 'cidrs' must be a list"
if len(cidrs) == 0:
return False, f"Site {i+1} must have at least one CIDR"
# Validate each CIDR
for j, cidr_config in enumerate(cidrs):
if not isinstance(cidr_config, dict):
return False, f"Site {i+1} CIDR {j+1} must be a dictionary/object"
if 'cidr' not in cidr_config:
return False, f"Site {i+1} CIDR {j+1} missing required field: 'cidr'"
# Validate CIDR format
cidr_str = cidr_config.get('cidr')
try:
ipaddress.ip_network(cidr_str, strict=False)
except ValueError:
return False, f"Site {i+1} CIDR {j+1}: Invalid CIDR notation '{cidr_str}'"
continue # Inline CIDR site is valid
# Legacy format: inline IPs
if 'name' not in site:
return False, f"Site {i+1} missing required field: 'name'"
if 'ips' not in site:
return False, f"Site {i+1} missing required field: 'ips'"
return False, f"Site {i+1} missing required field: 'ips' (or use 'site_ref' or 'cidrs')"
if not isinstance(site['ips'], list):
return False, f"Site {i+1} field 'ips' must be a list"
@@ -550,3 +927,60 @@ class ConfigService:
"""
filepath = os.path.join(self.configs_dir, filename)
return os.path.exists(filepath) and os.path.isfile(filepath)
def create_inline_sites(self, config_content: Dict) -> None:
"""
Create sites in the database for inline site definitions in a config.
This method scans the config for inline site definitions (with CIDRs)
and creates them as reusable sites in the database if they don't already exist.
Args:
config_content: Parsed YAML config dictionary
Raises:
ValueError: If site creation fails
"""
try:
from web.services.site_service import SiteService
from flask import current_app
site_service = SiteService(current_app.db_session)
sites = config_content.get('sites', [])
for site_def in sites:
# Skip site references (they already exist)
if 'site_ref' in site_def:
continue
# Skip legacy IP-based sites (not creating those as reusable sites)
if 'ips' in site_def and 'cidrs' not in site_def:
continue
# Process inline CIDR-based sites
if 'cidrs' in site_def:
site_name = site_def.get('name')
# Check if site already exists
existing_site = site_service.get_site_by_name(site_name)
if existing_site:
# Site already exists, skip creation
continue
# Create new site
cidrs = site_def.get('cidrs', [])
description = f"Auto-created from config '{config_content.get('title', 'Unknown')}'"
site_service.create_site(
name=site_name,
description=description,
cidrs=cidrs
)
except Exception as e:
# If site creation fails, log but don't block config creation
import logging
logging.getLogger(__name__).warning(
f"Failed to create inline sites from config: {str(e)}"
)

View File

@@ -16,7 +16,7 @@ from sqlalchemy.orm import Session, joinedload
from web.models import (
Scan, ScanSite, ScanIP, ScanPort, ScanService as ScanServiceModel,
ScanCertificate, ScanTLSVersion
ScanCertificate, ScanTLSVersion, Site, ScanSiteAssociation
)
from web.utils.pagination import paginate, PaginatedResult
from web.utils.validators import validate_config_file, validate_scan_status
@@ -41,8 +41,9 @@ class ScanService:
"""
self.db = db_session
def trigger_scan(self, config_file: str, triggered_by: str = 'manual',
schedule_id: Optional[int] = None, scheduler=None) -> int:
def trigger_scan(self, config_file: str = None, config_id: int = None,
triggered_by: str = 'manual', schedule_id: Optional[int] = None,
scheduler=None) -> int:
"""
Trigger a new scan.
@@ -50,7 +51,8 @@ class ScanService:
queues the scan for background execution.
Args:
config_file: Path to YAML configuration file
config_file: Path to YAML configuration file (legacy, optional)
config_id: Database config ID (preferred, optional)
triggered_by: Source that triggered scan (manual, scheduled, api)
schedule_id: Optional schedule ID if triggered by schedule
scheduler: Optional SchedulerService instance for queuing background jobs
@@ -59,57 +61,106 @@ class ScanService:
Scan ID of the created scan
Raises:
ValueError: If config file is invalid
ValueError: If config is invalid or both/neither config_file and config_id provided
"""
# Validate config file
is_valid, error_msg = validate_config_file(config_file)
if not is_valid:
raise ValueError(f"Invalid config file: {error_msg}")
# Validate that exactly one config source is provided
if not (bool(config_file) ^ bool(config_id)):
raise ValueError("Must provide exactly one of config_file or config_id")
# Convert config_file to full path if it's just a filename
if not config_file.startswith('/'):
config_path = f'/app/configs/{config_file}'
# Handle database config
if config_id:
from web.models import ScanConfig
# Validate config exists
db_config = self.db.query(ScanConfig).filter_by(id=config_id).first()
if not db_config:
raise ValueError(f"Config with ID {config_id} not found")
# Create scan record with config_id
scan = Scan(
timestamp=datetime.utcnow(),
status='running',
config_id=config_id,
title=db_config.title,
triggered_by=triggered_by,
schedule_id=schedule_id,
created_at=datetime.utcnow()
)
self.db.add(scan)
self.db.commit()
self.db.refresh(scan)
logger.info(f"Scan {scan.id} triggered via {triggered_by} with config_id={config_id}")
# Queue background job if scheduler provided
if scheduler:
try:
job_id = scheduler.queue_scan(scan.id, config_id=config_id)
logger.info(f"Scan {scan.id} queued for background execution (job_id={job_id})")
except Exception as e:
logger.error(f"Failed to queue scan {scan.id}: {str(e)}")
# Mark scan as failed if job queuing fails
scan.status = 'failed'
scan.error_message = f"Failed to queue background job: {str(e)}"
self.db.commit()
raise
else:
logger.warning(f"Scan {scan.id} created but not queued (no scheduler provided)")
return scan.id
# Handle legacy YAML config file
else:
config_path = config_file
# Validate config file
is_valid, error_msg = validate_config_file(config_file)
if not is_valid:
raise ValueError(f"Invalid config file: {error_msg}")
# Load config to get title
import yaml
with open(config_path, 'r') as f:
config = yaml.safe_load(f)
# Convert config_file to full path if it's just a filename
if not config_file.startswith('/'):
config_path = f'/app/configs/{config_file}'
else:
config_path = config_file
# Create scan record
scan = Scan(
timestamp=datetime.utcnow(),
status='running',
config_file=config_file,
title=config.get('title', 'Untitled Scan'),
triggered_by=triggered_by,
schedule_id=schedule_id,
created_at=datetime.utcnow()
)
# Load config to get title
import yaml
with open(config_path, 'r') as f:
config = yaml.safe_load(f)
self.db.add(scan)
self.db.commit()
self.db.refresh(scan)
# Create scan record
scan = Scan(
timestamp=datetime.utcnow(),
status='running',
config_file=config_file,
title=config.get('title', 'Untitled Scan'),
triggered_by=triggered_by,
schedule_id=schedule_id,
created_at=datetime.utcnow()
)
logger.info(f"Scan {scan.id} triggered via {triggered_by}")
self.db.add(scan)
self.db.commit()
self.db.refresh(scan)
# Queue background job if scheduler provided
if scheduler:
try:
job_id = scheduler.queue_scan(scan.id, config_file)
logger.info(f"Scan {scan.id} queued for background execution (job_id={job_id})")
except Exception as e:
logger.error(f"Failed to queue scan {scan.id}: {str(e)}")
# Mark scan as failed if job queuing fails
scan.status = 'failed'
scan.error_message = f"Failed to queue background job: {str(e)}"
self.db.commit()
raise
else:
logger.warning(f"Scan {scan.id} created but not queued (no scheduler provided)")
logger.info(f"Scan {scan.id} triggered via {triggered_by}")
return scan.id
# Queue background job if scheduler provided
if scheduler:
try:
job_id = scheduler.queue_scan(scan.id, config_file=config_file)
logger.info(f"Scan {scan.id} queued for background execution (job_id={job_id})")
except Exception as e:
logger.error(f"Failed to queue scan {scan.id}: {str(e)}")
# Mark scan as failed if job queuing fails
scan.status = 'failed'
scan.error_message = f"Failed to queue background job: {str(e)}"
self.db.commit()
raise
else:
logger.warning(f"Scan {scan.id} created but not queued (no scheduler provided)")
return scan.id
def get_scan(self, scan_id: int) -> Optional[Dict[str, Any]]:
"""
@@ -366,6 +417,34 @@ class ScanService:
self.db.add(site)
self.db.flush() # Get site.id for foreign key
# Create ScanSiteAssociation if this site exists in the database
# This links the scan to reusable site definitions
master_site = (
self.db.query(Site)
.filter(Site.name == site_data['name'])
.first()
)
if master_site:
# Check if association already exists (avoid duplicates)
existing_assoc = (
self.db.query(ScanSiteAssociation)
.filter(
ScanSiteAssociation.scan_id == scan_obj.id,
ScanSiteAssociation.site_id == master_site.id
)
.first()
)
if not existing_assoc:
assoc = ScanSiteAssociation(
scan_id=scan_obj.id,
site_id=master_site.id,
created_at=datetime.utcnow()
)
self.db.add(assoc)
logger.debug(f"Created association between scan {scan_obj.id} and site '{master_site.name}' (id={master_site.id})")
# Process each IP in this site
for ip_data in site_data.get('ips', []):
# Create ScanIP record

View File

@@ -149,13 +149,16 @@ class SchedulerService:
except Exception as e:
logger.error(f"Error loading schedules on startup: {str(e)}", exc_info=True)
def queue_scan(self, scan_id: int, config_file: str) -> str:
def queue_scan(self, scan_id: int, config_file: str = None, config_id: int = None) -> str:
"""
Queue a scan for immediate background execution.
Args:
scan_id: Database ID of the scan
config_file: Path to YAML configuration file
config_file: Path to YAML configuration file (legacy, optional)
config_id: Database config ID (preferred, optional)
Note: Provide exactly one of config_file or config_id
Returns:
Job ID from APScheduler
@@ -169,7 +172,7 @@ class SchedulerService:
# Add job to run immediately
job = self.scheduler.add_job(
func=execute_scan,
args=[scan_id, config_file, self.db_url],
kwargs={'scan_id': scan_id, 'config_file': config_file, 'config_id': config_id, 'db_url': self.db_url},
id=f'scan_{scan_id}',
name=f'Scan {scan_id}',
replace_existing=True,

View File

@@ -0,0 +1,531 @@
"""
Site service for managing reusable site definitions.
This service handles the business logic for creating, updating, and managing
sites with their associated CIDR ranges and IP-level overrides.
"""
import ipaddress
import json
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional
from sqlalchemy import func
from sqlalchemy.orm import Session, joinedload
from sqlalchemy.exc import IntegrityError
from web.models import (
Site, SiteCIDR, SiteIP, ScanSiteAssociation
)
from web.utils.pagination import paginate, PaginatedResult
logger = logging.getLogger(__name__)
class SiteService:
"""
Service for managing reusable site definitions.
Handles site lifecycle: creation, updates, deletion (with safety checks),
CIDR management, and IP-level overrides.
"""
def __init__(self, db_session: Session):
"""
Initialize site service.
Args:
db_session: SQLAlchemy database session
"""
self.db = db_session
def create_site(self, name: str, description: Optional[str] = None,
cidrs: Optional[List[Dict[str, Any]]] = None) -> Dict[str, Any]:
"""
Create a new site with optional CIDR ranges.
Args:
name: Unique site name
description: Optional site description
cidrs: List of CIDR definitions with format:
[{"cidr": "10.0.0.0/24", "expected_ping": true,
"expected_tcp_ports": [22, 80], "expected_udp_ports": [53]}]
Returns:
Dictionary with created site data
Raises:
ValueError: If site name already exists or validation fails
"""
# Validate site name is unique
existing = self.db.query(Site).filter(Site.name == name).first()
if existing:
raise ValueError(f"Site with name '{name}' already exists")
# Validate we have at least one CIDR if provided
if cidrs is not None and len(cidrs) == 0:
raise ValueError("Site must have at least one CIDR range")
# Create site
site = Site(
name=name,
description=description,
created_at=datetime.utcnow(),
updated_at=datetime.utcnow()
)
self.db.add(site)
self.db.flush() # Get site.id without committing
# Add CIDRs if provided
if cidrs:
for cidr_data in cidrs:
self._add_cidr_to_site(site, cidr_data)
self.db.commit()
self.db.refresh(site)
logger.info(f"Created site '{name}' (id={site.id}) with {len(cidrs or [])} CIDR(s)")
return self._site_to_dict(site)
def update_site(self, site_id: int, name: Optional[str] = None,
description: Optional[str] = None) -> Dict[str, Any]:
"""
Update site metadata (name and/or description).
Args:
site_id: Site ID to update
name: New site name (must be unique)
description: New description
Returns:
Dictionary with updated site data
Raises:
ValueError: If site not found or name already exists
"""
site = self.db.query(Site).filter(Site.id == site_id).first()
if not site:
raise ValueError(f"Site with id {site_id} not found")
# Update name if provided
if name is not None and name != site.name:
# Check uniqueness
existing = self.db.query(Site).filter(
Site.name == name,
Site.id != site_id
).first()
if existing:
raise ValueError(f"Site with name '{name}' already exists")
site.name = name
# Update description if provided
if description is not None:
site.description = description
site.updated_at = datetime.utcnow()
self.db.commit()
self.db.refresh(site)
logger.info(f"Updated site {site_id} ('{site.name}')")
return self._site_to_dict(site)
def delete_site(self, site_id: int) -> None:
"""
Delete a site.
Prevents deletion if the site is used in any scan (per user requirement).
Args:
site_id: Site ID to delete
Raises:
ValueError: If site not found or is used in scans
"""
site = self.db.query(Site).filter(Site.id == site_id).first()
if not site:
raise ValueError(f"Site with id {site_id} not found")
# Check if site is used in any scans
usage_count = (
self.db.query(func.count(ScanSiteAssociation.id))
.filter(ScanSiteAssociation.site_id == site_id)
.scalar()
)
if usage_count > 0:
raise ValueError(
f"Cannot delete site '{site.name}': it is used in {usage_count} scan(s). "
f"Sites that have been used in scans cannot be deleted."
)
# Safe to delete
self.db.delete(site)
self.db.commit()
logger.info(f"Deleted site {site_id} ('{site.name}')")
def get_site(self, site_id: int) -> Optional[Dict[str, Any]]:
"""
Get site details with all CIDRs and IP overrides.
Args:
site_id: Site ID to retrieve
Returns:
Dictionary with site data, or None if not found
"""
site = (
self.db.query(Site)
.options(
joinedload(Site.cidrs).joinedload(SiteCIDR.ips)
)
.filter(Site.id == site_id)
.first()
)
if not site:
return None
return self._site_to_dict(site)
def get_site_by_name(self, name: str) -> Optional[Dict[str, Any]]:
"""
Get site details by name.
Args:
name: Site name to retrieve
Returns:
Dictionary with site data, or None if not found
"""
site = (
self.db.query(Site)
.options(
joinedload(Site.cidrs).joinedload(SiteCIDR.ips)
)
.filter(Site.name == name)
.first()
)
if not site:
return None
return self._site_to_dict(site)
def list_sites(self, page: int = 1, per_page: int = 20) -> PaginatedResult:
"""
List all sites with pagination.
Args:
page: Page number (1-indexed)
per_page: Number of items per page
Returns:
PaginatedResult with site data
"""
query = (
self.db.query(Site)
.options(joinedload(Site.cidrs))
.order_by(Site.name)
)
return paginate(query, page, per_page, self._site_to_dict)
def list_all_sites(self) -> List[Dict[str, Any]]:
"""
List all sites without pagination (for dropdowns, etc.).
Returns:
List of site dictionaries
"""
sites = (
self.db.query(Site)
.options(joinedload(Site.cidrs))
.order_by(Site.name)
.all()
)
return [self._site_to_dict(site) for site in sites]
def add_cidr(self, site_id: int, cidr: str, expected_ping: Optional[bool] = None,
expected_tcp_ports: Optional[List[int]] = None,
expected_udp_ports: Optional[List[int]] = None) -> Dict[str, Any]:
"""
Add a CIDR range to a site.
Args:
site_id: Site ID
cidr: CIDR notation (e.g., "10.0.0.0/24")
expected_ping: Expected ping response for IPs in this CIDR
expected_tcp_ports: List of expected TCP ports
expected_udp_ports: List of expected UDP ports
Returns:
Dictionary with CIDR data
Raises:
ValueError: If site not found, CIDR is invalid, or already exists
"""
site = self.db.query(Site).filter(Site.id == site_id).first()
if not site:
raise ValueError(f"Site with id {site_id} not found")
# Validate CIDR format
try:
ipaddress.ip_network(cidr, strict=False)
except ValueError as e:
raise ValueError(f"Invalid CIDR notation '{cidr}': {str(e)}")
# Check for duplicate CIDR
existing = (
self.db.query(SiteCIDR)
.filter(SiteCIDR.site_id == site_id, SiteCIDR.cidr == cidr)
.first()
)
if existing:
raise ValueError(f"CIDR '{cidr}' already exists for this site")
# Create CIDR
cidr_obj = SiteCIDR(
site_id=site_id,
cidr=cidr,
expected_ping=expected_ping,
expected_tcp_ports=json.dumps(expected_tcp_ports or []),
expected_udp_ports=json.dumps(expected_udp_ports or []),
created_at=datetime.utcnow()
)
self.db.add(cidr_obj)
site.updated_at = datetime.utcnow()
self.db.commit()
self.db.refresh(cidr_obj)
logger.info(f"Added CIDR '{cidr}' to site {site_id} ('{site.name}')")
return self._cidr_to_dict(cidr_obj)
def remove_cidr(self, site_id: int, cidr_id: int) -> None:
"""
Remove a CIDR range from a site.
Prevents removal if it's the last CIDR (sites must have at least one CIDR).
Args:
site_id: Site ID
cidr_id: CIDR ID to remove
Raises:
ValueError: If CIDR not found or it's the last CIDR
"""
site = self.db.query(Site).filter(Site.id == site_id).first()
if not site:
raise ValueError(f"Site with id {site_id} not found")
cidr = (
self.db.query(SiteCIDR)
.filter(SiteCIDR.id == cidr_id, SiteCIDR.site_id == site_id)
.first()
)
if not cidr:
raise ValueError(f"CIDR with id {cidr_id} not found for site {site_id}")
# Check if this is the last CIDR
cidr_count = (
self.db.query(func.count(SiteCIDR.id))
.filter(SiteCIDR.site_id == site_id)
.scalar()
)
if cidr_count <= 1:
raise ValueError(
f"Cannot remove CIDR '{cidr.cidr}': site must have at least one CIDR range"
)
self.db.delete(cidr)
site.updated_at = datetime.utcnow()
self.db.commit()
logger.info(f"Removed CIDR '{cidr.cidr}' from site {site_id} ('{site.name}')")
def add_ip_override(self, cidr_id: int, ip_address: str,
expected_ping: Optional[bool] = None,
expected_tcp_ports: Optional[List[int]] = None,
expected_udp_ports: Optional[List[int]] = None) -> Dict[str, Any]:
"""
Add an IP-level expectation override within a CIDR.
Args:
cidr_id: CIDR ID
ip_address: IP address to override
expected_ping: Override ping expectation
expected_tcp_ports: Override TCP ports expectation
expected_udp_ports: Override UDP ports expectation
Returns:
Dictionary with IP override data
Raises:
ValueError: If CIDR not found, IP is invalid, or not in CIDR range
"""
cidr = self.db.query(SiteCIDR).filter(SiteCIDR.id == cidr_id).first()
if not cidr:
raise ValueError(f"CIDR with id {cidr_id} not found")
# Validate IP format
try:
ip_obj = ipaddress.ip_address(ip_address)
except ValueError as e:
raise ValueError(f"Invalid IP address '{ip_address}': {str(e)}")
# Validate IP is within CIDR range
network = ipaddress.ip_network(cidr.cidr, strict=False)
if ip_obj not in network:
raise ValueError(f"IP address '{ip_address}' is not within CIDR '{cidr.cidr}'")
# Check for duplicate
existing = (
self.db.query(SiteIP)
.filter(SiteIP.site_cidr_id == cidr_id, SiteIP.ip_address == ip_address)
.first()
)
if existing:
raise ValueError(f"IP override for '{ip_address}' already exists in this CIDR")
# Create IP override
ip_override = SiteIP(
site_cidr_id=cidr_id,
ip_address=ip_address,
expected_ping=expected_ping,
expected_tcp_ports=json.dumps(expected_tcp_ports or []),
expected_udp_ports=json.dumps(expected_udp_ports or []),
created_at=datetime.utcnow()
)
self.db.add(ip_override)
self.db.commit()
self.db.refresh(ip_override)
logger.info(f"Added IP override '{ip_address}' to CIDR {cidr_id} ('{cidr.cidr}')")
return self._ip_override_to_dict(ip_override)
def remove_ip_override(self, cidr_id: int, ip_id: int) -> None:
"""
Remove an IP-level override.
Args:
cidr_id: CIDR ID
ip_id: IP override ID to remove
Raises:
ValueError: If IP override not found
"""
ip_override = (
self.db.query(SiteIP)
.filter(SiteIP.id == ip_id, SiteIP.site_cidr_id == cidr_id)
.first()
)
if not ip_override:
raise ValueError(f"IP override with id {ip_id} not found for CIDR {cidr_id}")
ip_address = ip_override.ip_address
self.db.delete(ip_override)
self.db.commit()
logger.info(f"Removed IP override '{ip_address}' from CIDR {cidr_id}")
def get_scan_usage(self, site_id: int) -> List[Dict[str, Any]]:
"""
Get list of scans that use this site.
Args:
site_id: Site ID
Returns:
List of scan dictionaries
"""
from web.models import Scan # Import here to avoid circular dependency
associations = (
self.db.query(ScanSiteAssociation)
.options(joinedload(ScanSiteAssociation.scan))
.filter(ScanSiteAssociation.site_id == site_id)
.all()
)
return [
{
'id': assoc.scan.id,
'title': assoc.scan.title,
'timestamp': assoc.scan.timestamp.isoformat() if assoc.scan.timestamp else None,
'status': assoc.scan.status
}
for assoc in associations
]
# Private helper methods
def _add_cidr_to_site(self, site: Site, cidr_data: Dict[str, Any]) -> SiteCIDR:
"""Helper to add CIDR during site creation."""
cidr = cidr_data.get('cidr')
if not cidr:
raise ValueError("CIDR 'cidr' field is required")
# Validate CIDR format
try:
ipaddress.ip_network(cidr, strict=False)
except ValueError as e:
raise ValueError(f"Invalid CIDR notation '{cidr}': {str(e)}")
cidr_obj = SiteCIDR(
site_id=site.id,
cidr=cidr,
expected_ping=cidr_data.get('expected_ping'),
expected_tcp_ports=json.dumps(cidr_data.get('expected_tcp_ports', [])),
expected_udp_ports=json.dumps(cidr_data.get('expected_udp_ports', [])),
created_at=datetime.utcnow()
)
self.db.add(cidr_obj)
return cidr_obj
def _site_to_dict(self, site: Site) -> Dict[str, Any]:
"""Convert Site model to dictionary."""
return {
'id': site.id,
'name': site.name,
'description': site.description,
'created_at': site.created_at.isoformat() if site.created_at else None,
'updated_at': site.updated_at.isoformat() if site.updated_at else None,
'cidrs': [self._cidr_to_dict(cidr) for cidr in site.cidrs] if hasattr(site, 'cidrs') else []
}
def _cidr_to_dict(self, cidr: SiteCIDR) -> Dict[str, Any]:
"""Convert SiteCIDR model to dictionary."""
return {
'id': cidr.id,
'site_id': cidr.site_id,
'cidr': cidr.cidr,
'expected_ping': cidr.expected_ping,
'expected_tcp_ports': json.loads(cidr.expected_tcp_ports) if cidr.expected_tcp_ports else [],
'expected_udp_ports': json.loads(cidr.expected_udp_ports) if cidr.expected_udp_ports else [],
'created_at': cidr.created_at.isoformat() if cidr.created_at else None,
'ip_overrides': [self._ip_override_to_dict(ip) for ip in cidr.ips] if hasattr(cidr, 'ips') else []
}
def _ip_override_to_dict(self, ip: SiteIP) -> Dict[str, Any]:
"""Convert SiteIP model to dictionary."""
return {
'id': ip.id,
'site_cidr_id': ip.site_cidr_id,
'ip_address': ip.ip_address,
'expected_ping': ip.expected_ping,
'expected_tcp_ports': json.loads(ip.expected_tcp_ports) if ip.expected_tcp_ports else [],
'expected_udp_ports': json.loads(ip.expected_udp_ports) if ip.expected_udp_ports else [],
'created_at': ip.created_at.isoformat() if ip.created_at else None
}