Kitos
51c927394d
- Replace default=datetime.utcnow with server_default=func.now() across all 16 models (17 columns) for consistent, timezone-aware timestamps from PostgreSQL - Upgrade DateTime columns to DateTime(timezone=True) for timestamptz storage - Configure SQLAlchemy engine pool: pool_size=20, max_overflow=10, pool_recycle=3600, pool_pre_ping=True - Remove unused datetime imports from model files
153 lines
4.7 KiB
Python
153 lines
4.7 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
|
|
from sqlalchemy import (
|
|
Column, String, Text, Integer, Boolean, DateTime,
|
|
ForeignKey, Index, func,
|
|
)
|
|
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
|
from sqlalchemy.orm import relationship
|
|
|
|
from app.database import Base
|
|
|
|
|
|
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
|
|
"""
|
|
__tablename__ = "campaigns"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
name = Column(String, nullable=False)
|
|
description = Column(Text, nullable=True)
|
|
type = Column(String, nullable=False, default="custom") # custom, apt_emulation, kill_chain, compliance
|
|
threat_actor_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("threat_actors.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
)
|
|
status = Column(String, nullable=False, default="draft") # draft, active, completed, archived
|
|
created_by = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("users.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
)
|
|
scheduled_at = Column(DateTime, nullable=True)
|
|
completed_at = Column(DateTime, nullable=True)
|
|
target_platform = Column(String, nullable=True)
|
|
tags = Column(JSONB, nullable=True, default=[])
|
|
created_at = Column(DateTime(timezone=True), server_default=func.now())
|
|
|
|
# Recurring scheduling fields
|
|
is_recurring = Column(Boolean, default=False)
|
|
recurrence_pattern = Column(String, nullable=True) # weekly, monthly, quarterly
|
|
next_run_at = Column(DateTime, nullable=True)
|
|
last_run_at = Column(DateTime, nullable=True)
|
|
parent_campaign_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("campaigns.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
)
|
|
|
|
# Relationships
|
|
threat_actor = relationship("ThreatActor")
|
|
creator = relationship("User", foreign_keys=[created_by])
|
|
campaign_tests = relationship(
|
|
"CampaignTest",
|
|
back_populates="campaign",
|
|
cascade="all, delete-orphan",
|
|
order_by="CampaignTest.order_index",
|
|
)
|
|
parent_campaign = relationship(
|
|
"Campaign",
|
|
remote_side="Campaign.id",
|
|
foreign_keys=[parent_campaign_id],
|
|
)
|
|
child_campaigns = relationship(
|
|
"Campaign",
|
|
foreign_keys=[parent_campaign_id],
|
|
back_populates="parent_campaign",
|
|
)
|
|
|
|
__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 = [
|
|
"reconnaissance",
|
|
"resource_development",
|
|
"initial_access",
|
|
"execution",
|
|
"persistence",
|
|
"privilege_escalation",
|
|
"defense_evasion",
|
|
"credential_access",
|
|
"discovery",
|
|
"lateral_movement",
|
|
"collection",
|
|
"command_and_control",
|
|
"exfiltration",
|
|
"impact",
|
|
]
|
|
|
|
|
|
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.
|
|
"""
|
|
__tablename__ = "campaign_tests"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
campaign_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("campaigns.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
)
|
|
test_id = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("tests.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
)
|
|
order_index = Column(Integer, nullable=False, default=0)
|
|
depends_on = Column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("campaign_tests.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
)
|
|
phase = Column(String, nullable=True) # kill chain phase
|
|
|
|
# Relationships
|
|
campaign = relationship("Campaign", back_populates="campaign_tests")
|
|
test = relationship("Test")
|
|
dependency = relationship("CampaignTest", remote_side="CampaignTest.id")
|
|
|
|
__table_args__ = (
|
|
Index('ix_campaign_tests_campaign', 'campaign_id'),
|
|
Index('ix_campaign_tests_test', 'test_id'),
|
|
)
|