feat(refactor): PEP8, type annotations, docstrings and PyJWT security fix
This commit is contained in:
kitos
@@ -4,20 +4,35 @@ 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 (
|
||||
Column, String, Text, Integer, Boolean, DateTime,
|
||||
ForeignKey, Index, func,
|
||||
Boolean,
|
||||
Column,
|
||||
DateTime,
|
||||
ForeignKey,
|
||||
Index,
|
||||
Integer,
|
||||
String,
|
||||
Text,
|
||||
func,
|
||||
)
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
|
||||
# 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.
|
||||
"""A campaign groups multiple tests into a sequenced attack chain.
|
||||
|
||||
Types:
|
||||
- custom: manually created campaign
|
||||
@@ -31,62 +46,97 @@ class Campaign(Base):
|
||||
- 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,
|
||||
)
|
||||
start_date = Column(DateTime, nullable=True) # campaign won't activate before this date
|
||||
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'),
|
||||
@@ -98,56 +148,83 @@ class Campaign(Base):
|
||||
|
||||
# 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.
|
||||
"""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'),
|
||||
|
||||
Reference in New Issue
Block a user