diff --git a/controler/main.py b/controler/main.py index a3bb1ce..96527ce 100644 --- a/controler/main.py +++ b/controler/main.py @@ -3,12 +3,10 @@ API Controller for A/D Infrastructure Manages docker-compose services with authentication """ import os -import subprocess import asyncio from datetime import datetime -from typing import Optional, List +from typing import Optional, List, Tuple from fastapi import FastAPI, HTTPException, Depends, Header -from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from pydantic import BaseModel import asyncpg from contextlib import asynccontextmanager @@ -27,14 +25,11 @@ class ServiceCreate(BaseModel): git_url: Optional[str] = None class ServiceAction(BaseModel): - action: str # start, stop, restart + action: str class GitPullRequest(BaseModel): auto: bool = False -class LogRequest(BaseModel): - lines: int = 100 - # Auth dependency async def verify_token(authorization: str = Header(None)): if not authorization or not authorization.startswith("Bearer "): @@ -114,7 +109,7 @@ def parse_docker_status(ps_output: str) -> dict: 'containers': containers } -async def run_docker_compose_command(service_path: str, command: List[str]) -> tuple[int, str, str]: +async def run_docker_compose_command(service_path: str, command: List[str]) -> Tuple[int, str, str]: """Run docker-compose command and return (returncode, stdout, stderr)""" compose_file = find_compose_file(service_path) if not compose_file: diff --git a/scoreboard_injector/main.py b/scoreboard_injector/main.py index 99877d5..d46ed35 100644 --- a/scoreboard_injector/main.py +++ b/scoreboard_injector/main.py @@ -3,13 +3,11 @@ Scoreboard Injector for ForcAD Monitors Socket.IO events for attacks and alerts on critical situations """ import os -import json import asyncio +import aiohttp from datetime import datetime, timedelta -from typing import Optional, Dict, Any import socketio from fastapi import FastAPI, HTTPException, Depends, Header -from pydantic import BaseModel import asyncpg from contextlib import asynccontextmanager @@ -19,20 +17,12 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production") SCOREBOARD_URL = os.getenv("SCOREBOARD_URL", "http://10.60.0.1:8080") OUR_TEAM_ID = int(os.getenv("OUR_TEAM_ID", "1")) ALERT_THRESHOLD_POINTS = float(os.getenv("ALERT_THRESHOLD_POINTS", "5")) -ALERT_THRESHOLD_TIME = int(os.getenv("ALERT_THRESHOLD_TIME", "300")) # seconds TELEGRAM_API_URL = os.getenv("TELEGRAM_API_URL", "http://tg-bot:8003/send") # Database pool db_pool = None ws_task = None -class AttackStats(BaseModel): - total_attacks: int - attacks_by_us: int - attacks_to_us: int - recent_attacks: int - critical_alerts: int - # Auth dependency async def verify_token(authorization: str = Header(None)): if not authorization or not authorization.startswith("Bearer "): @@ -43,168 +33,36 @@ async def verify_token(authorization: str = Header(None)): raise HTTPException(status_code=403, detail="Invalid token") return token -# Database functions -async def get_db(): - return await db_pool.acquire() -async def release_db(conn): - await db_pool.release(conn) - -async def process_attack_event(event: Dict[str, Any]): - """Process attack event from scoreboard""" - conn = await db_pool.acquire() - try: - # Extract attack information from event - # Handle multiple possible event formats from ForcAD - event_type = event.get('type', 'unknown') - - # Try to extract attacker/victim IDs from various possible fields - attacker_id = event.get('attacker_id') or event.get('team_id') or event.get('attacker') - victim_id = event.get('victim_id') or event.get('target_id') or event.get('victim') or event.get('target') - - # Skip if we don't have both attacker and victim - if attacker_id is None or victim_id is None: - print(f"Skipping event with missing attacker/victim: {event}") - return - - # Convert to integers if they're strings - try: - attacker_id = int(attacker_id) - victim_id = int(victim_id) - except (ValueError, TypeError): - print(f"Invalid team IDs in event: attacker={attacker_id}, victim={victim_id}") - return - - service_name = event.get('service') or event.get('service_name') or event.get('task_name') or 'unknown' - flag = event.get('flag', '') - - # Handle timestamp - time_str = event.get('time') or event.get('timestamp') - if time_str: - try: - # Try parsing ISO format - timestamp = datetime.fromisoformat(time_str.replace('Z', '+00:00')) - except (ValueError, AttributeError): - # Try parsing as Unix timestamp - try: - timestamp = datetime.fromtimestamp(float(time_str)) - except (ValueError, TypeError): - timestamp = datetime.utcnow() - else: - timestamp = datetime.utcnow() - - # Extract points (might be in different fields) - points = float(event.get('points', 0) or event.get('score', 0) or 1.0) - - # Generate unique attack ID - round_num = event.get('round', event.get('round_id', 0)) - attack_id = event.get('id') or f"{round_num}_{attacker_id}_{victim_id}_{service_name}_{int(timestamp.timestamp())}" - - is_our_attack = attacker_id == OUR_TEAM_ID - is_attack_to_us = victim_id == OUR_TEAM_ID - - # Only log if it involves our team - if is_our_attack or is_attack_to_us: - # Store attack in database - inserted = await conn.fetchval(""" - INSERT INTO attacks (attack_id, attacker_team_id, victim_team_id, service_name, flag, timestamp, points, is_our_attack, is_attack_to_us) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) - ON CONFLICT (attack_id) DO NOTHING - RETURNING id - """, attack_id, attacker_id, victim_id, service_name, flag, timestamp, points, is_our_attack, is_attack_to_us) - - if inserted: - print(f"[{event_type}] Logged attack: Team {attacker_id} -> Team {victim_id} | {service_name} | {points} pts") - - # Check for alert conditions if attack is against us - if is_attack_to_us: - await check_and_create_alerts(conn, attacker_id, service_name) - - except Exception as e: - print(f"Error processing attack event: {e}") - print(f"Event data: {event}") - finally: - await db_pool.release(conn) - -async def check_and_create_alerts(conn, attacker_id: int, service_name: str): - """Check if we should create an alert for attacks against us""" - threshold_time = datetime.utcnow() - timedelta(seconds=ALERT_THRESHOLD_TIME) - - # Check total points lost from this attacker in threshold time - result = await conn.fetchrow(""" - SELECT COUNT(*) as attack_count, COALESCE(SUM(points), 0) as total_points - FROM attacks - WHERE is_attack_to_us = true - AND attacker_team_id = $1 - AND service_name = $2 - AND timestamp > $3 - """, attacker_id, service_name, threshold_time) - - if result and result['total_points'] >= ALERT_THRESHOLD_POINTS: - # Create alert - alert_message = f"CRITICAL: Team {attacker_id} has stolen {result['total_points']:.2f} points from service {service_name} in the last {ALERT_THRESHOLD_TIME}s ({result['attack_count']} attacks)" - - # Check if we already alerted recently - recent_alert = await conn.fetchrow(""" - SELECT id FROM attack_alerts - WHERE alert_type = 'high_point_loss' - AND message LIKE $1 - AND created_at > $2 - """, f"%Team {attacker_id}%{service_name}%", threshold_time) - - if not recent_alert: - alert_id = await conn.fetchval(""" - INSERT INTO attack_alerts (attack_id, alert_type, severity, message) - VALUES ( - (SELECT id FROM attacks WHERE attacker_team_id = $1 AND service_name = $2 ORDER BY timestamp DESC LIMIT 1), - 'high_point_loss', - 'critical', - $3 - ) - RETURNING id - """, attacker_id, service_name, alert_message) - - # Send to telegram - await send_telegram_alert(alert_message) - - # Mark as notified - await conn.execute("UPDATE attack_alerts SET notified = true WHERE id = $1", alert_id) - -async def send_telegram_alert(message: str): +async def send_telegram_alert(message: str, service_id: int = None, service_name: str = None): """Send alert to telegram bot""" - import aiohttp try: - print(f"📱 Sending alert to Telegram: {TELEGRAM_API_URL}") async with aiohttp.ClientSession() as session: + payload = {"message": message} + if service_id: + payload["service_id"] = service_id + if service_name: + payload["service_name"] = service_name + async with session.post( TELEGRAM_API_URL, - json={"message": message}, + json=payload, headers={"Authorization": f"Bearer {SECRET_TOKEN}"} ) as resp: - response_text = await resp.text() if resp.status != 200: - print(f"❌ Failed to send telegram alert: Status {resp.status}") - print(f" Response: {response_text}") - else: - print(f"✅ Telegram alert sent successfully") - print(f" Response: {response_text}") + print(f"Failed to send telegram alert: Status {resp.status}") except Exception as e: - print(f"❌ Error sending telegram alert: {e}") - import traceback - traceback.print_exc() + print(f"Error sending telegram alert: {e}") async def fetch_task_names(): """Fetch task names from scoreboard API""" - import aiohttp try: async with aiohttp.ClientSession() as session: async with session.get(f"{SCOREBOARD_URL}/api/client/tasks/") as resp: if resp.status == 200: tasks = await resp.json() return {task['id']: task['name'] for task in tasks} - else: - print(f"Failed to fetch tasks: {resp.status}") - return {} + return {} except Exception as e: print(f"Error fetching task names: {e}") return {} @@ -219,44 +77,35 @@ async def socketio_listener(): # Fetch task names on startup task_names.update(await fetch_task_names()) - if task_names: - print(f"📋 Loaded task names: {', '.join([f'{name} (ID:{tid})' for tid, name in task_names.items()])}") @sio.on('*', namespace='/live_events') async def catch_all(event, data): """Catch all events from live_events namespace""" - print(f"📡 Received event: {event}") - print(f" Data: {data}") - - # Parse the event format: ["event_type", {"data": ...}] if isinstance(data, list) and len(data) >= 2: event_type = data[0] event_data = data[1].get('data', {}) if isinstance(data[1], dict) else {} if event_type == 'flag_stolen': await process_flag_stolen(event_data) - elif isinstance(data, dict): - # Handle direct event data - if 'data' in data: - await process_flag_stolen(data['data']) + elif isinstance(data, dict) and 'data' in data: + await process_flag_stolen(data['data']) async def process_flag_stolen(event_data): """Process flag_stolen event""" try: - print(f"[DEBUG] process_flag_stolen called with event_data: {event_data}") attacker_id = event_data.get('attacker_id') victim_id = event_data.get('victim_id') task_id = event_data.get('task_id') attacker_delta = event_data.get('attacker_delta', 0) - print(f"[DEBUG] attacker_id={attacker_id}, victim_id={victim_id}, task_id={task_id}, attacker_delta={attacker_delta}") + if attacker_id is None or victim_id is None: - print("[DEBUG] attacker_id or victim_id is None, skipping event") return + service_name = task_names.get(task_id, f"task_{task_id}") timestamp = datetime.utcnow() is_our_attack = attacker_id == OUR_TEAM_ID is_attack_to_us = victim_id == OUR_TEAM_ID - print(f"[DEBUG] is_our_attack={is_our_attack}, is_attack_to_us={is_attack_to_us}, ALERT_THRESHOLD_POINTS={ALERT_THRESHOLD_POINTS}") + if is_our_attack or is_attack_to_us: conn = await db_pool.acquire() try: @@ -266,34 +115,38 @@ async def socketio_listener(): VALUES ($1, $2, $3, $4, $5, $6, $7, $8) ON CONFLICT (attack_id) DO NOTHING """, attack_id, attacker_id, victim_id, service_name, timestamp, float(attacker_delta), is_our_attack, is_attack_to_us) - if is_our_attack: - print(f" ✅ We stole flag from Team {victim_id} on {service_name} (+{attacker_delta:.2f} FP)") - elif is_attack_to_us: - print(f" ⚠ī¸ Team {attacker_id} stole flag from us on {service_name} (-{attacker_delta:.2f} FP)") - if attacker_delta >= ALERT_THRESHOLD_POINTS: - print(f"[DEBUG] Sending Telegram alert: attacker_delta={attacker_delta} >= ALERT_THRESHOLD_POINTS={ALERT_THRESHOLD_POINTS}") - alert_message = f"🚨 ATTACK DETECTED!\nTeam {attacker_id} stole flag from {service_name}\nPoints lost: {attacker_delta:.2f} FP" - alert_id = await conn.fetchval(""" - INSERT INTO attack_alerts (attack_id, alert_type, severity, message) - VALUES ( - (SELECT id FROM attacks WHERE attack_id = $1), - 'flag_stolen', - 'high', - $2 - ) - RETURNING id - """, attack_id, alert_message) - await send_telegram_alert(alert_message) - await conn.execute("UPDATE attack_alerts SET notified = true WHERE id = $1", alert_id) - print(f" 📱 Alert sent to Telegram") - else: - print(f"[DEBUG] No alert sent: attacker_delta={attacker_delta} < ALERT_THRESHOLD_POINTS={ALERT_THRESHOLD_POINTS}") + + if is_attack_to_us and attacker_delta >= ALERT_THRESHOLD_POINTS: + alert_message = f"🚨 ATTACK DETECTED!\nTeam {attacker_id} stole flag from {service_name}\nPoints lost: {attacker_delta:.2f} FP" + + # Get service_id from controller if available + service_id = None + try: + service_row = await conn.fetchrow( + "SELECT id FROM services WHERE name = $1 LIMIT 1", + service_name + ) + if service_row: + service_id = service_row['id'] + except: + pass + + alert_id = await conn.fetchval(""" + INSERT INTO attack_alerts (attack_id, alert_type, severity, message) + VALUES ( + (SELECT id FROM attacks WHERE attack_id = $1), + 'flag_stolen', + 'high', + $2 + ) + RETURNING id + """, attack_id, alert_message) + await send_telegram_alert(alert_message, service_id=service_id, service_name=service_name) + await conn.execute("UPDATE attack_alerts SET notified = true WHERE id = $1", alert_id) finally: await db_pool.release(conn) except Exception as e: print(f"Error processing flag_stolen event: {e}") - import traceback - traceback.print_exc() @sio.event(namespace='/live_events') async def update_scoreboard(data): @@ -304,8 +157,6 @@ async def socketio_listener(): round_start = event_data.get('round_start', 0) team_tasks = event_data.get('team_tasks', []) - print(f"📊 Round {round_num} - Processing {len(team_tasks)} team updates") - conn = await db_pool.acquire() try: # Store team scores from team_tasks (score field = FP for this service) @@ -408,59 +259,43 @@ async def socketio_listener(): """, attack_id, attacker_id, victim_id, service_name, timestamp, float(fp_value), is_our_attack, is_attack_to_us) if is_our_attack: - print(f" ✅ We stole {new_stolen} flags from {service_name} (+{fp_value:.2f} FP)") + pass elif is_attack_to_us: - print(f" ⚠ī¸ We LOST {new_lost} flags on {service_name} (-{fp_value:.2f} FP)") if fp_value >= ALERT_THRESHOLD_POINTS: await check_and_create_alerts(conn, 0, service_name) - elif new_stolen > 0: - print(f" 📌 Team {team_id} stole {new_stolen} flags from {service_name} (+{fp_change:.2f} FP)") - elif new_lost > 0: - print(f" 📌 Team {team_id} lost {new_lost} flags on {service_name} ({fp_change:.2f} FP)") - finally: + finally: await db_pool.release(conn) except Exception as e: print(f"Error processing update_scoreboard: {e}") - import traceback - traceback.print_exc() @sio.event(namespace='/live_events') async def init_scoreboard(data): """Handle initial scoreboard data""" try: - print("📡 Received initial scoreboard data") event_data = data.get('data', {}) teams = event_data.get('teams', []) tasks = event_data.get('tasks', []) - # Cache task names for task in tasks: task_names[task.get('id')] = task.get('name') - # Cache team names for team in teams: team_names[team.get('id')] = team.get('name') - - team_names_str = ', '.join([f"{t.get('name')} (ID:{t.get('id')})" for t in teams]) - task_names_str = ', '.join([t.get('name') for t in tasks]) - print(f" Teams: {team_names_str}") - print(f" Tasks: {task_names_str}") except Exception as e: print(f"Error processing init_scoreboard: {e}") @sio.event async def connect(): - print("✅ Connected to ForcAD scoreboard Socket.IO") + pass @sio.event async def disconnect(): - print("❌ Disconnected from scoreboard") + pass while True: try: - print(f"Connecting to {SCOREBOARD_URL}/socket.io ...") await sio.connect( SCOREBOARD_URL, namespaces=['/live_events'], @@ -468,8 +303,6 @@ async def socketio_listener(): ) await sio.wait() except Exception as e: - print(f"Socket.IO error: {e}") - print("Reconnecting in 5 seconds...") await asyncio.sleep(5) # Lifespan context @@ -477,16 +310,10 @@ async def socketio_listener(): async def lifespan(app: FastAPI): global db_pool, ws_task db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10) - - print(f"Starting Socket.IO listener") - print(f"Scoreboard URL: {SCOREBOARD_URL}") - print(f"Our team ID: {OUR_TEAM_ID}") - ws_task = asyncio.create_task(socketio_listener()) yield - # Cleanup if ws_task: ws_task.cancel() try: diff --git a/tg-bot/main.py b/tg-bot/main.py index af9a5f5..19fb411 100644 --- a/tg-bot/main.py +++ b/tg-bot/main.py @@ -7,8 +7,9 @@ from datetime import datetime from fastapi import FastAPI, HTTPException, Depends, Header from pydantic import BaseModel import asyncpg -from telegram import Bot +from telegram import Bot, InlineKeyboardButton, InlineKeyboardMarkup, Update from telegram.error import TelegramError +from telegram.ext import Application, CallbackQueryHandler, ContextTypes from contextlib import asynccontextmanager # Configuration @@ -17,19 +18,56 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production") TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "") TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID", "") -if not TELEGRAM_BOT_TOKEN: - print("WARNING: TELEGRAM_BOT_TOKEN not set!") - -if not TELEGRAM_CHAT_ID: - print("WARNING: TELEGRAM_CHAT_ID not set!") - # Database pool and bot db_pool = None bot = None +app_telegram = None + +async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE): + """Handle inline button clicks""" + query = update.callback_query + await query.answer() + + callback_data = query.data + chat_id = query.message.chat_id + + if callback_data.startswith("service_"): + action, service_id = callback_data.rsplit("_", 1) + action = action.replace("service_", "") + + try: + import aiohttp + controller_url = os.getenv("CONTROLLER_API", "http://controller:8001") + + async with aiohttp.ClientSession() as session: + api_url = f"{controller_url}/services/{service_id}/action" + headers = {"Authorization": f"Bearer {SECRET_TOKEN}"} + data = {"action": action} + + async with session.post(api_url, json=data, headers=headers) as resp: + if resp.status == 200: + result = await resp.json() + await query.edit_message_text( + text=f"✅ Service action '{action}' executed successfully\n{query.message.text}" + ) + await log_message(chat_id, f"Button action: {action} service {service_id}", True) + else: + error_text = await resp.text() + await query.edit_message_text( + text=f"❌ Failed to execute action: {error_text}\n{query.message.text}" + ) + await log_message(chat_id, f"Button action failed: {action} service {service_id}", False, error_text) + except Exception as e: + await query.edit_message_text( + text=f"❌ Error: {str(e)}\n{query.message.text}" + ) + await log_message(chat_id, f"Button action error: {callback_data}", False, str(e)) class MessageRequest(BaseModel): message: str - chat_id: str = None # Optional, uses default if not provided + chat_id: str = None + service_id: int = None + service_name: str = None class BulkMessageRequest(BaseModel): messages: list[str] @@ -66,14 +104,22 @@ async def log_message(chat_id: int, message: str, success: bool, error_message: # Lifespan context @asynccontextmanager async def lifespan(app: FastAPI): - global db_pool, bot + global db_pool, bot, app_telegram db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10) if TELEGRAM_BOT_TOKEN: bot = Bot(token=TELEGRAM_BOT_TOKEN) + + app_telegram = Application.builder().token(TELEGRAM_BOT_TOKEN).build() + app_telegram.add_handler(CallbackQueryHandler(button_handler)) + + await app_telegram.initialize() yield + if app_telegram: + await app_telegram.shutdown() + await db_pool.close() app = FastAPI(title="Telegram Bot API", lifespan=lifespan) @@ -89,7 +135,7 @@ async def health_check(): @app.post("/send", dependencies=[Depends(verify_token)]) async def send_message(request: MessageRequest): - """Send a message to telegram chat""" + """Send a message to telegram chat with optional service control buttons""" if not bot: raise HTTPException(status_code=503, detail="Telegram bot not configured") @@ -98,14 +144,24 @@ async def send_message(request: MessageRequest): raise HTTPException(status_code=400, detail="No chat_id provided and no default configured") try: - # Send message - message = await bot.send_message( - chat_id=int(chat_id), - text=request.message, - parse_mode='HTML' - ) + kwargs = { + "chat_id": int(chat_id), + "text": request.message, + "parse_mode": "HTML" + } - # Log success + # Add inline buttons for service control if service_id is provided + if request.service_id and request.service_name: + keyboard = [ + [ + InlineKeyboardButton("â–ļī¸ Start", callback_data=f"service_start_{request.service_id}"), + InlineKeyboardButton("⏚ī¸ Stop", callback_data=f"service_stop_{request.service_id}"), + InlineKeyboardButton("🔄 Restart", callback_data=f"service_restart_{request.service_id}") + ] + ] + kwargs["reply_markup"] = InlineKeyboardMarkup(keyboard) + + message = await bot.send_message(**kwargs) await log_message(int(chat_id), request.message, True) return { @@ -115,7 +171,6 @@ async def send_message(request: MessageRequest): } except TelegramError as e: - # Log failure await log_message(int(chat_id), request.message, False, str(e)) raise HTTPException(status_code=500, detail=f"Failed to send message: {str(e)}") @@ -184,6 +239,20 @@ async def get_stats(): finally: await release_db(conn) +@app.post("/webhook") +async def webhook(update_data: dict): + """Handle Telegram webhook updates for button callbacks""" + if not app_telegram: + raise HTTPException(status_code=503, detail="Telegram app not configured") + + try: + update = Update.de_json(update_data, bot) + if update: + await app_telegram.process_update(update) + return {"status": "ok"} + except Exception as e: + return {"status": "error", "message": str(e)} + @app.post("/test", dependencies=[Depends(verify_token)]) async def test_connection(): """Test telegram bot connection""" diff --git a/tg-bot/requirements.txt b/tg-bot/requirements.txt index 361db73..9073141 100644 --- a/tg-bot/requirements.txt +++ b/tg-bot/requirements.txt @@ -2,5 +2,6 @@ fastapi==0.109.0 uvicorn[standard]==0.27.0 asyncpg==0.29.0 pydantic==2.5.3 -python-telegram-bot==21.0 +python-telegram-bot[all]==21.0 python-dotenv==1.0.0 +aiohttp==3.9.1 diff --git a/web/app.py b/web/app.py index 0108ec0..b844415 100644 --- a/web/app.py +++ b/web/app.py @@ -4,14 +4,12 @@ Flask-based dashboard to monitor services, attacks, and alerts """ import os import asyncio -from datetime import datetime, timedelta +from datetime import datetime from flask import Flask, render_template, jsonify, request, redirect, url_for, session -import asyncpg import aiohttp from functools import wraps # Configuration -DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://adctrl:adctrl@postgres:5432/adctrl") SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production") WEB_PASSWORD = os.getenv("WEB_PASSWORD", "admin123") CONTROLLER_API = os.getenv("CONTROLLER_API", "http://controller:8001") @@ -21,10 +19,6 @@ TELEGRAM_API = os.getenv("TELEGRAM_API", "http://tg-bot:8003") app = Flask(__name__) app.secret_key = os.getenv("FLASK_SECRET_KEY", "change-me-in-production-flask-secret") -# Database connection -async def get_db_conn(): - return await asyncpg.connect(DATABASE_URL) - # Auth decorator def login_required(f): @wraps(f)