"""D3FEND query service — framework-agnostic queries for defensive techniques.""" # Enable future language features for compatibility from __future__ import annotations # Import Optional from typing from typing import Optional # Import func from sqlalchemy from sqlalchemy import func # Import Session from sqlalchemy.orm from sqlalchemy.orm import Session # Import EntityNotFoundError from app.domain.errors from app.domain.errors import EntityNotFoundError # Import DefensiveTechnique from app.models.defensive_technique from app.models.defensive_technique import DefensiveTechnique # Import Technique from app.models.technique from app.models.technique import Technique # Import get_defenses_for_technique from app.services.d3fend_import_service from app.services.d3fend_import_service import get_defenses_for_technique # Import escape_like from app.utils from app.utils import escape_like # Define function list_defensive_techniques def list_defensive_techniques( # Entry: db db: Session, *, # Entry: tactic tactic: Optional[str] = None, # Entry: search search: Optional[str] = None, # Entry: offset offset: int = 0, # Entry: limit limit: int = 50, ) -> dict: """List D3FEND defensive techniques with optional filters.""" # Assign query = db.query(DefensiveTechnique) query = db.query(DefensiveTechnique) # Check: tactic if tactic: # Assign query = query.filter(DefensiveTechnique.tactic == tactic) query = query.filter(DefensiveTechnique.tactic == tactic) # Check: search if search: # Assign pattern = f"%{escape_like(search)}%" pattern = f"%{escape_like(search)}%" # Assign query = query.filter( query = query.filter( DefensiveTechnique.name.ilike(pattern) | DefensiveTechnique.d3fend_id.ilike(pattern) ) # Assign total = query.count() total = query.count() # Assign items = query.order_by(DefensiveTechnique.d3fend_id).offset(offset).limit(l... items = query.order_by(DefensiveTechnique.d3fend_id).offset(offset).limit(limit).all() # Return { return { # Literal argument value "total": total, # Literal argument value "offset": offset, # Literal argument value "limit": limit, # Literal argument value "items": [ { # Literal argument value "id": str(dt.id), # Literal argument value "d3fend_id": dt.d3fend_id, # Literal argument value "name": dt.name, # Literal argument value "description": dt.description, # Literal argument value "tactic": dt.tactic, # Literal argument value "d3fend_url": dt.d3fend_url, } for dt in items ], } # Define function list_d3fend_tactics def list_d3fend_tactics(db: Session) -> list[dict]: """Return a list of all D3FEND tactics with counts.""" # Assign rows = ( rows = ( db.query(DefensiveTechnique.tactic, func.count(DefensiveTechnique.id)) # Chain .group_by() call .group_by(DefensiveTechnique.tactic) # Chain .order_by() call .order_by(DefensiveTechnique.tactic) # Chain .all() call .all() ) # Return [{"tactic": tactic or "Unknown", "count": count} for tactic, count ... return [{"tactic": tactic or "Unknown", "count": count} for tactic, count in rows] # Define function get_defenses_for_attack_technique def get_defenses_for_attack_technique(db: Session, mitre_id: str) -> dict: """Get all D3FEND defensive techniques mapped to a given ATT&CK technique.""" # Assign technique = db.query(Technique).filter(Technique.mitre_id == mitre_id).first() technique = db.query(Technique).filter(Technique.mitre_id == mitre_id).first() # Check: technique is None if technique is None: # Raise EntityNotFoundError raise EntityNotFoundError("Technique", mitre_id) # Assign defenses = get_defenses_for_technique(db, technique.id) defenses = get_defenses_for_technique(db, technique.id) # Return { return { # Literal argument value "mitre_id": mitre_id, # Literal argument value "technique_name": technique.name, # Literal argument value "defenses": defenses, # Literal argument value "total": len(defenses), }