feat(backend): c2 crypto + config CRUD + adapter scaffolding (sprint 8 M1)
- Add Fernet crypto service (MIMIC_ENCRYPTION_KEY env, C2Disabled on absent key) - Add Alembic migration 0006: c2_config + c2_task tables with cascade FKs - Add C2Config and C2Task SQLAlchemy models - Add C2Adapter ABC with dataclasses (C2Health, C2Callback, C2TaskStatus, C2TaskPage) - Add FakeAdapter (deterministic in-memory, MIMIC_C2_ADAPTER=fake) - Add MythicAdapter scaffold: test_connection() live, M2+ raise NotImplementedError - Add decode_response_text() helper for base64/binary Mythic responses - Add GET/PUT/DELETE/POST-test /api/engagements/<id>/c2-config endpoints - RBAC: admin+redteam OK, SOC 403; 503 guard when encryption key absent - Token never returned in API responses; stored Fernet-encrypted only - 42 new tests (300 total, 258 baseline preserved green) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Knacky
30
backend/tests/test_c2_adapter_fake.py
Normal file
30
backend/tests/test_c2_adapter_fake.py
Normal file
@@ -0,0 +1,30 @@
|
||||
"""Tests for the FakeAdapter deterministic in-memory implementation."""
|
||||
from __future__ import annotations
|
||||
|
||||
from backend.app.services.c2.adapter import C2Health
|
||||
from backend.app.services.c2.fake import FakeAdapter
|
||||
|
||||
|
||||
class TestFakeAdapterTestConnection:
|
||||
def test_returns_ok_true(self):
|
||||
adapter = FakeAdapter()
|
||||
health = adapter.test_connection()
|
||||
assert isinstance(health, C2Health)
|
||||
assert health.ok is True
|
||||
assert health.error is None
|
||||
|
||||
def test_list_callbacks_returns_list(self):
|
||||
adapter = FakeAdapter()
|
||||
callbacks = adapter.list_callbacks()
|
||||
assert isinstance(callbacks, list)
|
||||
assert len(callbacks) >= 1
|
||||
|
||||
def test_list_callbacks_fields(self):
|
||||
adapter = FakeAdapter()
|
||||
cb = adapter.list_callbacks()[0]
|
||||
assert hasattr(cb, "display_id")
|
||||
assert hasattr(cb, "active")
|
||||
assert hasattr(cb, "host")
|
||||
assert hasattr(cb, "user")
|
||||
assert hasattr(cb, "domain")
|
||||
assert hasattr(cb, "last_checkin")
|
||||
352
backend/tests/test_c2_config.py
Normal file
352
backend/tests/test_c2_config.py
Normal file
@@ -0,0 +1,352 @@
|
||||
"""Tests for C2 config CRUD endpoints.
|
||||
|
||||
Covers:
|
||||
- GET 404 when no config exists
|
||||
- PUT create (api_token required)
|
||||
- PUT update with omitted token keeps old ciphertext
|
||||
- GET 200 returns has_token=True, never cleartext
|
||||
- DELETE 204
|
||||
- Cascade delete when engagement is deleted
|
||||
- RBAC: admin OK / redteam OK / SOC 403 on all 4 endpoints
|
||||
- 503 guard when MIMIC_ENCRYPTION_KEY is unset
|
||||
- POST /test with fake adapter
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
from flask import Flask
|
||||
from flask.testing import FlaskClient
|
||||
|
||||
from backend.app.models.c2_config import C2Config
|
||||
from backend.tests.conftest import auth_headers as _h
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_FERNET_KEY = Fernet.generate_key().decode()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def set_encryption_key(monkeypatch):
|
||||
"""Default: key is present. Individual tests can override."""
|
||||
monkeypatch.setenv("MIMIC_ENCRYPTION_KEY", _FERNET_KEY)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_fake_adapter(monkeypatch):
|
||||
monkeypatch.setenv("MIMIC_C2_ADAPTER", "fake")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_engagement(client: FlaskClient, token: str) -> dict:
|
||||
resp = client.post(
|
||||
"/api/engagements",
|
||||
headers=_h(token),
|
||||
json={"name": "Op Alpha", "start_date": "2026-06-10"},
|
||||
)
|
||||
assert resp.status_code == 201, resp.get_json()
|
||||
return resp.get_json()
|
||||
|
||||
|
||||
def _put_config(
|
||||
client: FlaskClient,
|
||||
token: str,
|
||||
eid: int,
|
||||
*,
|
||||
url: str = "https://c2.internal:7443",
|
||||
api_token: str | None = "s3cr3t",
|
||||
verify_tls: bool = True,
|
||||
) -> dict:
|
||||
payload: dict = {"url": url, "verify_tls": verify_tls}
|
||||
if api_token is not None:
|
||||
payload["api_token"] = api_token
|
||||
resp = client.put(
|
||||
f"/api/engagements/{eid}/c2-config",
|
||||
headers=_h(token),
|
||||
json=payload,
|
||||
)
|
||||
return resp
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET — 404 when no config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_config_not_found(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
resp = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token))
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
def test_get_config_engagement_not_found(client: FlaskClient, admin_token: str) -> None:
|
||||
resp = client.get("/api/engagements/9999/c2-config", headers=_h(admin_token))
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PUT — create
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_put_creates_config(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
resp = _put_config(client, admin_token, eng["id"])
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
assert body["has_token"] is True
|
||||
assert body["url"] == "https://c2.internal:7443"
|
||||
assert body["verify_tls"] is True
|
||||
|
||||
|
||||
def test_put_create_requires_api_token(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
resp = _put_config(client, admin_token, eng["id"], api_token=None)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_put_create_requires_url(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
resp = client.put(
|
||||
f"/api/engagements/{eng['id']}/c2-config",
|
||||
headers=_h(admin_token),
|
||||
json={"api_token": "tok", "verify_tls": True},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PUT — update, omitting api_token preserves old ciphertext
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_put_update_omits_token_keeps_old(
|
||||
app: Flask, client: FlaskClient, admin_token: str
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
_put_config(client, admin_token, eng["id"], api_token="original-token")
|
||||
|
||||
# Read ciphertext from DB before update.
|
||||
with app.app_context():
|
||||
cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first()
|
||||
assert cfg is not None
|
||||
old_cipher = cfg.api_token_encrypted
|
||||
|
||||
# Update URL, omit api_token.
|
||||
resp = _put_config(
|
||||
client, admin_token, eng["id"],
|
||||
url="https://new.internal:7443", api_token=None,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
with app.app_context():
|
||||
cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first()
|
||||
assert cfg is not None
|
||||
assert cfg.api_token_encrypted == old_cipher
|
||||
assert cfg.url == "https://new.internal:7443"
|
||||
|
||||
|
||||
def test_put_update_with_token_replaces_ciphertext(
|
||||
app: Flask, client: FlaskClient, admin_token: str
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
_put_config(client, admin_token, eng["id"], api_token="original-token")
|
||||
|
||||
with app.app_context():
|
||||
cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first()
|
||||
assert cfg is not None
|
||||
old_cipher = cfg.api_token_encrypted
|
||||
|
||||
_put_config(client, admin_token, eng["id"], api_token="new-token")
|
||||
|
||||
with app.app_context():
|
||||
cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first()
|
||||
assert cfg is not None
|
||||
assert cfg.api_token_encrypted != old_cipher
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET — 200, has_token=True, never cleartext
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_config_returns_has_token_not_cleartext(
|
||||
client: FlaskClient, admin_token: str
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
_put_config(client, admin_token, eng["id"], api_token="s3cr3t")
|
||||
|
||||
resp = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token))
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
assert body["has_token"] is True
|
||||
assert "api_token" not in body
|
||||
assert "api_token_encrypted" not in body
|
||||
assert "s3cr3t" not in str(body)
|
||||
|
||||
|
||||
def test_get_config_verify_tls_default_true(
|
||||
client: FlaskClient, admin_token: str
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
_put_config(client, admin_token, eng["id"])
|
||||
resp = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token))
|
||||
assert resp.get_json()["verify_tls"] is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DELETE — 204
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delete_config_204(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
_put_config(client, admin_token, eng["id"])
|
||||
|
||||
resp = client.delete(
|
||||
f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token)
|
||||
)
|
||||
assert resp.status_code == 204
|
||||
|
||||
# Subsequent GET returns 404.
|
||||
resp2 = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token))
|
||||
assert resp2.status_code == 404
|
||||
|
||||
|
||||
def test_delete_config_not_found(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
resp = client.delete(
|
||||
f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token)
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CASCADE — delete engagement removes config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_cascade_delete_engagement_removes_config(
|
||||
app: Flask, client: FlaskClient, admin_token: str
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
_put_config(client, admin_token, eng["id"])
|
||||
|
||||
with app.app_context():
|
||||
assert C2Config.query.filter_by(engagement_id=eng["id"]).count() == 1
|
||||
|
||||
client.delete(f"/api/engagements/{eng['id']}", headers=_h(admin_token))
|
||||
|
||||
with app.app_context():
|
||||
assert C2Config.query.filter_by(engagement_id=eng["id"]).count() == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RBAC matrix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method,path_suffix", [
|
||||
("GET", "/c2-config"),
|
||||
("PUT", "/c2-config"),
|
||||
("DELETE", "/c2-config"),
|
||||
("POST", "/c2-config/test"),
|
||||
])
|
||||
def test_soc_gets_403(
|
||||
client: FlaskClient, admin_token: str, soc_token: str,
|
||||
method: str, path_suffix: str,
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
url = f"/api/engagements/{eng['id']}{path_suffix}"
|
||||
resp = getattr(client, method.lower())(url, headers=_h(soc_token), json={})
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method,path_suffix", [
|
||||
("GET", "/c2-config"),
|
||||
("DELETE", "/c2-config"),
|
||||
("POST", "/c2-config/test"),
|
||||
])
|
||||
def test_redteam_gets_allowed(
|
||||
client: FlaskClient, admin_token: str, redteam_token: str,
|
||||
method: str, path_suffix: str,
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
# Ensure config exists for GET/DELETE/test.
|
||||
_put_config(client, admin_token, eng["id"])
|
||||
url = f"/api/engagements/{eng['id']}{path_suffix}"
|
||||
resp = getattr(client, method.lower())(url, headers=_h(redteam_token), json={})
|
||||
# Not 403 and not 401.
|
||||
assert resp.status_code not in (401, 403)
|
||||
|
||||
|
||||
def test_redteam_can_put_config(
|
||||
client: FlaskClient, admin_token: str, redteam_token: str,
|
||||
) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
resp = _put_config(client, redteam_token, eng["id"])
|
||||
assert resp.status_code == 200
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 503 guard when MIMIC_ENCRYPTION_KEY is unset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method,path_suffix", [
|
||||
("GET", "/c2-config"),
|
||||
("PUT", "/c2-config"),
|
||||
("DELETE", "/c2-config"),
|
||||
("POST", "/c2-config/test"),
|
||||
])
|
||||
def test_503_when_key_unset(
|
||||
monkeypatch,
|
||||
client: FlaskClient,
|
||||
admin_token: str,
|
||||
method: str,
|
||||
path_suffix: str,
|
||||
) -> None:
|
||||
monkeypatch.delenv("MIMIC_ENCRYPTION_KEY", raising=False)
|
||||
eng = _make_engagement(client, admin_token)
|
||||
url = f"/api/engagements/{eng['id']}{path_suffix}"
|
||||
resp = getattr(client, method.lower())(url, headers=_h(admin_token), json={
|
||||
"url": "https://c2", "api_token": "tok", "verify_tls": True,
|
||||
})
|
||||
assert resp.status_code == 503
|
||||
assert "MIMIC_ENCRYPTION_KEY" in resp.get_json().get("error", "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST /test — connectivity check via fake adapter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_post_test_returns_ok_true(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
_put_config(client, admin_token, eng["id"])
|
||||
|
||||
resp = client.post(
|
||||
f"/api/engagements/{eng['id']}/c2-config/test",
|
||||
headers=_h(admin_token),
|
||||
json={},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
assert body["ok"] is True
|
||||
assert body["error"] is None
|
||||
|
||||
|
||||
def test_post_test_no_config_returns_404(client: FlaskClient, admin_token: str) -> None:
|
||||
eng = _make_engagement(client, admin_token)
|
||||
resp = client.post(
|
||||
f"/api/engagements/{eng['id']}/c2-config/test",
|
||||
headers=_h(admin_token),
|
||||
json={},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
52
backend/tests/test_crypto.py
Normal file
52
backend/tests/test_crypto.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""Tests for the Fernet crypto service."""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from backend.app.services.crypto import C2Disabled, decrypt, encrypt
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fernet_key(monkeypatch) -> str:
|
||||
key = Fernet.generate_key().decode()
|
||||
monkeypatch.setenv("MIMIC_ENCRYPTION_KEY", key)
|
||||
return key
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def no_key(monkeypatch):
|
||||
monkeypatch.delenv("MIMIC_ENCRYPTION_KEY", raising=False)
|
||||
|
||||
|
||||
class TestEncryptDecrypt:
|
||||
def test_round_trip(self, fernet_key):
|
||||
plaintext = "s3cr3t-api-token"
|
||||
ciphertext = encrypt(plaintext)
|
||||
assert ciphertext != plaintext
|
||||
assert decrypt(ciphertext) == plaintext
|
||||
|
||||
def test_different_tokens_for_same_input(self, fernet_key):
|
||||
# Fernet tokens are non-deterministic (random IV).
|
||||
t1 = encrypt("same")
|
||||
t2 = encrypt("same")
|
||||
assert t1 != t2
|
||||
assert decrypt(t1) == decrypt(t2) == "same"
|
||||
|
||||
def test_decrypt_invalid_ciphertext(self, fernet_key):
|
||||
with pytest.raises(ValueError):
|
||||
decrypt("not-valid-fernet-token")
|
||||
|
||||
|
||||
class TestKeyAbsent:
|
||||
def test_encrypt_raises_c2disabled(self, no_key):
|
||||
with pytest.raises(C2Disabled):
|
||||
encrypt("anything")
|
||||
|
||||
def test_decrypt_raises_c2disabled(self, no_key):
|
||||
with pytest.raises(C2Disabled):
|
||||
decrypt("anything")
|
||||
|
||||
def test_c2disabled_message(self, no_key):
|
||||
with pytest.raises(C2Disabled, match="MIMIC_ENCRYPTION_KEY"):
|
||||
encrypt("x")
|
||||
204
backend/tests/test_migration_0006_c2.py
Normal file
204
backend/tests/test_migration_0006_c2.py
Normal file
@@ -0,0 +1,204 @@
|
||||
"""Migration round-trip test for 0006_c2_layer.
|
||||
|
||||
Verifies that upgrade() creates c2_config and c2_task with the expected schema,
|
||||
and that downgrade() removes both tables cleanly.
|
||||
|
||||
Uses the resolved-path pattern (derives path from __file__) to avoid the
|
||||
hardcoded-path regression documented in lessons.md Sprint 4.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
from alembic.operations import Operations
|
||||
from alembic.runtime.migration import MigrationContext
|
||||
from sqlalchemy import create_engine, inspect, text
|
||||
|
||||
|
||||
def _load_migration():
|
||||
versions_dir = Path(__file__).resolve().parent.parent / "migrations" / "versions"
|
||||
path = versions_dir / "0006_c2_layer.py"
|
||||
spec = importlib.util.spec_from_file_location("migration_0006", path)
|
||||
assert spec and spec.loader
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod) # type: ignore[union-attr]
|
||||
return mod
|
||||
|
||||
|
||||
def _fresh_engine():
|
||||
"""In-memory SQLite with the tables that 0006 depends on already present."""
|
||||
engine = create_engine("sqlite:///:memory:")
|
||||
with engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
CREATE TABLE users (
|
||||
id INTEGER PRIMARY KEY,
|
||||
username TEXT NOT NULL UNIQUE,
|
||||
password_hash TEXT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
CREATE TABLE engagements (
|
||||
id INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
start_date DATE NOT NULL,
|
||||
end_date DATE,
|
||||
status TEXT NOT NULL DEFAULT 'planned',
|
||||
created_at DATETIME NOT NULL,
|
||||
created_by_id INTEGER NOT NULL REFERENCES users(id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
CREATE TABLE simulations (
|
||||
id INTEGER PRIMARY KEY,
|
||||
engagement_id INTEGER NOT NULL REFERENCES engagements(id),
|
||||
name TEXT NOT NULL,
|
||||
techniques JSON NOT NULL DEFAULT '[]',
|
||||
tactic_ids JSON NOT NULL DEFAULT '[]',
|
||||
description TEXT,
|
||||
commands TEXT,
|
||||
prerequisites TEXT,
|
||||
executed_at DATETIME,
|
||||
execution_result TEXT,
|
||||
log_source TEXT,
|
||||
logs TEXT,
|
||||
soc_comment TEXT,
|
||||
incident_number TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME,
|
||||
created_by_id INTEGER NOT NULL REFERENCES users(id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
return engine
|
||||
|
||||
|
||||
def _run_upgrade(engine, migration_mod):
|
||||
with engine.begin() as conn:
|
||||
ctx = MigrationContext.configure(conn)
|
||||
ops = Operations(ctx)
|
||||
# Patch op module for the migration
|
||||
import alembic.op as alembic_op
|
||||
original_proxy = alembic_op._proxy # type: ignore[attr-defined]
|
||||
alembic_op._proxy = ops # type: ignore[attr-defined]
|
||||
try:
|
||||
migration_mod.upgrade()
|
||||
finally:
|
||||
alembic_op._proxy = original_proxy # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _run_downgrade(engine, migration_mod):
|
||||
with engine.begin() as conn:
|
||||
ctx = MigrationContext.configure(conn)
|
||||
ops = Operations(ctx)
|
||||
import alembic.op as alembic_op
|
||||
original_proxy = alembic_op._proxy # type: ignore[attr-defined]
|
||||
alembic_op._proxy = ops # type: ignore[attr-defined]
|
||||
try:
|
||||
migration_mod.downgrade()
|
||||
finally:
|
||||
alembic_op._proxy = original_proxy # type: ignore[attr-defined]
|
||||
|
||||
|
||||
class TestMigration0006Upgrade:
|
||||
def test_c2_config_table_created(self):
|
||||
engine = _fresh_engine()
|
||||
mod = _load_migration()
|
||||
_run_upgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
assert "c2_config" in insp.get_table_names()
|
||||
|
||||
def test_c2_task_table_created(self):
|
||||
engine = _fresh_engine()
|
||||
mod = _load_migration()
|
||||
_run_upgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
assert "c2_task" in insp.get_table_names()
|
||||
|
||||
def test_c2_config_columns(self):
|
||||
engine = _fresh_engine()
|
||||
mod = _load_migration()
|
||||
_run_upgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
cols = {c["name"] for c in insp.get_columns("c2_config")}
|
||||
assert {"id", "engagement_id", "url", "api_token_encrypted",
|
||||
"verify_tls", "created_at", "updated_at"} <= cols
|
||||
|
||||
def test_c2_config_unique_constraint_on_engagement_id(self):
|
||||
engine = _fresh_engine()
|
||||
mod = _load_migration()
|
||||
_run_upgrade(engine, mod)
|
||||
|
||||
# Insert a user and engagement first.
|
||||
with engine.begin() as conn:
|
||||
conn.execute(text(
|
||||
"INSERT INTO users (id, username, password_hash, role, created_at) "
|
||||
"VALUES (1, 'u', 'h', 'admin', '2026-01-01')"
|
||||
))
|
||||
conn.execute(text(
|
||||
"INSERT INTO engagements (id, name, start_date, status, created_at, created_by_id) "
|
||||
"VALUES (1, 'Op', '2026-01-01', 'planned', '2026-01-01', 1)"
|
||||
))
|
||||
conn.execute(text(
|
||||
"INSERT INTO c2_config (engagement_id, url, api_token_encrypted, verify_tls, created_at) "
|
||||
"VALUES (1, 'https://c2', 'tok', 1, '2026-01-01')"
|
||||
))
|
||||
# Second insert on same engagement_id must fail.
|
||||
try:
|
||||
conn.execute(text(
|
||||
"INSERT INTO c2_config (engagement_id, url, api_token_encrypted, verify_tls, created_at) "
|
||||
"VALUES (1, 'https://c2b', 'tok2', 1, '2026-01-01')"
|
||||
))
|
||||
raised = False
|
||||
except Exception:
|
||||
raised = True
|
||||
assert raised, "UNIQUE constraint on c2_config.engagement_id must be enforced"
|
||||
|
||||
def test_c2_task_columns(self):
|
||||
engine = _fresh_engine()
|
||||
mod = _load_migration()
|
||||
_run_upgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
cols = {c["name"] for c in insp.get_columns("c2_task")}
|
||||
assert {"id", "simulation_id", "mythic_task_display_id", "callback_display_id",
|
||||
"command", "params", "status", "completed", "output", "source",
|
||||
"created_at", "completed_at"} <= cols
|
||||
|
||||
|
||||
class TestMigration0006Downgrade:
|
||||
def test_downgrade_removes_c2_config(self):
|
||||
engine = _fresh_engine()
|
||||
mod = _load_migration()
|
||||
_run_upgrade(engine, mod)
|
||||
_run_downgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
assert "c2_config" not in insp.get_table_names()
|
||||
|
||||
def test_downgrade_removes_c2_task(self):
|
||||
engine = _fresh_engine()
|
||||
mod = _load_migration()
|
||||
_run_upgrade(engine, mod)
|
||||
_run_downgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
assert "c2_task" not in insp.get_table_names()
|
||||
Reference in New Issue
Block a user