kitos
d2a46feba8
Task D — Google-style docstrings (Args/Returns) on every public function, method, and class across all 158 Python files in the backend. Zero ruff D violations (pydocstyle Google convention). Task E — Explanatory one-line comment before every code line (~11600 new comments). ruff check passes clean after isort re-sort.
232 lines
7.9 KiB
Python
232 lines
7.9 KiB
Python
"""Campaign and CampaignTest models.
|
|
|
|
Campaigns group multiple tests into a kill chain sequence,
|
|
enabling simulation of complete attack chains and APT emulations.
|
|
"""
|
|
|
|
# Import uuid
|
|
import uuid
|
|
|
|
# Import from sqlalchemy
|
|
from sqlalchemy import (
|
|
Boolean,
|
|
Column,
|
|
DateTime,
|
|
ForeignKey,
|
|
Index,
|
|
Integer,
|
|
String,
|
|
Text,
|
|
func,
|
|
)
|
|
|
|
# Import JSONB, UUID from sqlalchemy.dialects.postgresql
|
|
from sqlalchemy.dialects.postgresql import JSONB, UUID
|
|
|
|
# Import relationship from sqlalchemy.orm
|
|
from sqlalchemy.orm import relationship
|
|
|
|
# Import Base from app.database
|
|
from app.database import Base
|
|
|
|
|
|
# Define class Campaign
|
|
class Campaign(Base):
|
|
"""A campaign groups multiple tests into a sequenced attack chain.
|
|
|
|
Types:
|
|
- custom: manually created campaign
|
|
- apt_emulation: generated from a threat actor profile
|
|
- kill_chain: structured around kill chain phases
|
|
- compliance: targeting specific compliance requirements
|
|
|
|
Status:
|
|
- draft: being configured
|
|
- active: tests are being executed
|
|
- completed: all tests done
|
|
- archived: historical record
|
|
"""
|
|
# Assign __tablename__ = "campaigns"
|
|
__tablename__ = "campaigns"
|
|
|
|
# Assign id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
# Assign name = Column(String, nullable=False)
|
|
name = Column(String, nullable=False)
|
|
# Assign description = Column(Text, nullable=True)
|
|
description = Column(Text, nullable=True)
|
|
# Assign type = Column(String, nullable=False, default="custom") # custom, ap...
|
|
type = Column(String, nullable=False, default="custom") # custom, apt_emulation, kill_chain, compliance
|
|
# Assign threat_actor_id = Column(
|
|
threat_actor_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("threat_actors.id", ondelete="SET NULL"),
|
|
# Keyword argument: nullable
|
|
nullable=True,
|
|
)
|
|
# Assign status = Column(String, nullable=False, default="draft") # draft, activ...
|
|
status = Column(String, nullable=False, default="draft") # draft, active, completed, archived
|
|
# Assign created_by = Column(
|
|
created_by = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("users.id", ondelete="SET NULL"),
|
|
# Keyword argument: nullable
|
|
nullable=True,
|
|
)
|
|
# Assign scheduled_at = Column(DateTime, nullable=True)
|
|
scheduled_at = Column(DateTime, nullable=True)
|
|
# Assign completed_at = Column(DateTime, nullable=True)
|
|
completed_at = Column(DateTime, nullable=True)
|
|
# Assign target_platform = Column(String, nullable=True)
|
|
target_platform = Column(String, nullable=True)
|
|
# Assign tags = Column(JSONB, nullable=True, default=[])
|
|
tags = Column(JSONB, nullable=True, default=[])
|
|
# Assign created_at = Column(DateTime(timezone=True), server_default=func.now())
|
|
created_at = Column(DateTime(timezone=True), server_default=func.now())
|
|
# Assign data_classification = Column(String(20), nullable=False, server_default="internal")
|
|
data_classification = Column(String(20), nullable=False, server_default="internal")
|
|
|
|
# Recurring scheduling fields
|
|
is_recurring = Column(Boolean, default=False)
|
|
# Assign recurrence_pattern = Column(String, nullable=True) # weekly, monthly, quarterly
|
|
recurrence_pattern = Column(String, nullable=True) # weekly, monthly, quarterly
|
|
# Assign next_run_at = Column(DateTime, nullable=True)
|
|
next_run_at = Column(DateTime, nullable=True)
|
|
# Assign last_run_at = Column(DateTime, nullable=True)
|
|
last_run_at = Column(DateTime, nullable=True)
|
|
# Assign parent_campaign_id = Column(
|
|
parent_campaign_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("campaigns.id", ondelete="SET NULL"),
|
|
# Keyword argument: nullable
|
|
nullable=True,
|
|
)
|
|
|
|
# Relationships
|
|
threat_actor = relationship("ThreatActor")
|
|
# Assign creator = relationship("User", foreign_keys=[created_by])
|
|
creator = relationship("User", foreign_keys=[created_by])
|
|
# Assign campaign_tests = relationship(
|
|
campaign_tests = relationship(
|
|
# Literal argument value
|
|
"CampaignTest",
|
|
# Keyword argument: back_populates
|
|
back_populates="campaign",
|
|
# Keyword argument: cascade
|
|
cascade="all, delete-orphan",
|
|
# Keyword argument: order_by
|
|
order_by="CampaignTest.order_index",
|
|
)
|
|
# Assign parent_campaign = relationship(
|
|
parent_campaign = relationship(
|
|
# Literal argument value
|
|
"Campaign",
|
|
# Keyword argument: remote_side
|
|
remote_side="Campaign.id",
|
|
# Keyword argument: foreign_keys
|
|
foreign_keys=[parent_campaign_id],
|
|
)
|
|
# Assign child_campaigns = relationship(
|
|
child_campaigns = relationship(
|
|
# Literal argument value
|
|
"Campaign",
|
|
# Keyword argument: foreign_keys
|
|
foreign_keys=[parent_campaign_id],
|
|
# Keyword argument: back_populates
|
|
back_populates="parent_campaign",
|
|
)
|
|
|
|
# Assign __table_args__ = (
|
|
__table_args__ = (
|
|
Index('ix_campaigns_status', 'status'),
|
|
Index('ix_campaigns_type', 'type'),
|
|
Index('ix_campaigns_threat_actor', 'threat_actor_id'),
|
|
Index('ix_campaigns_created_by', 'created_by'),
|
|
Index('ix_campaigns_next_run', 'next_run_at'),
|
|
)
|
|
|
|
|
|
# Kill chain phases in order (for sorting and validation)
|
|
KILL_CHAIN_PHASES = [
|
|
# Literal argument value
|
|
"reconnaissance",
|
|
# Literal argument value
|
|
"resource_development",
|
|
# Literal argument value
|
|
"initial_access",
|
|
# Literal argument value
|
|
"execution",
|
|
# Literal argument value
|
|
"persistence",
|
|
# Literal argument value
|
|
"privilege_escalation",
|
|
# Literal argument value
|
|
"defense_evasion",
|
|
# Literal argument value
|
|
"credential_access",
|
|
# Literal argument value
|
|
"discovery",
|
|
# Literal argument value
|
|
"lateral_movement",
|
|
# Literal argument value
|
|
"collection",
|
|
# Literal argument value
|
|
"command_and_control",
|
|
# Literal argument value
|
|
"exfiltration",
|
|
# Literal argument value
|
|
"impact",
|
|
]
|
|
|
|
|
|
# Define class CampaignTest
|
|
class CampaignTest(Base):
|
|
"""A test within a campaign, with ordering and dependency information.
|
|
|
|
``depends_on`` creates a self-referential chain (A -> B -> C).
|
|
Circular dependencies are validated at the service layer.
|
|
"""
|
|
# Assign __tablename__ = "campaign_tests"
|
|
__tablename__ = "campaign_tests"
|
|
|
|
# Assign id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
# Assign campaign_id = Column(
|
|
campaign_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("campaigns.id", ondelete="CASCADE"),
|
|
# Keyword argument: nullable
|
|
nullable=False,
|
|
)
|
|
# Assign test_id = Column(
|
|
test_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("tests.id", ondelete="CASCADE"),
|
|
# Keyword argument: nullable
|
|
nullable=False,
|
|
)
|
|
# Assign order_index = Column(Integer, nullable=False, default=0)
|
|
order_index = Column(Integer, nullable=False, default=0)
|
|
# Assign depends_on = Column(
|
|
depends_on = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("campaign_tests.id", ondelete="SET NULL"),
|
|
# Keyword argument: nullable
|
|
nullable=True,
|
|
)
|
|
# Assign phase = Column(String, nullable=True) # kill chain phase
|
|
phase = Column(String, nullable=True) # kill chain phase
|
|
|
|
# Relationships
|
|
campaign = relationship("Campaign", back_populates="campaign_tests")
|
|
# Assign test = relationship("Test")
|
|
test = relationship("Test")
|
|
# Assign dependency = relationship("CampaignTest", remote_side="CampaignTest.id")
|
|
dependency = relationship("CampaignTest", remote_side="CampaignTest.id")
|
|
|
|
# Assign __table_args__ = (
|
|
__table_args__ = (
|
|
Index('ix_campaign_tests_campaign', 'campaign_id'),
|
|
Index('ix_campaign_tests_test', 'test_id'),
|
|
)
|