0xb00b5/ad-infr-control
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ads

Browse Source
This commit is contained in:
ilyastar9999
2025-12-04 13:26:39 +03:00
parent 34662c2a11
commit 7f841d155d
5 changed files with 142 additions and 256 deletions
Download Patch File Download Diff File Expand all files Collapse all files

271
scoreboard_injector/main.py
View File

@@ -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: