Files
ad-infr-control/scoreboard_injector/main.py
ilyastar9999 d69ccb038c asd
2025-12-04 14:09:19 +03:00

584 lines
25 KiB
Python

"""
Scoreboard Injector for ForcAD
Monitors Socket.IO events for attacks and alerts on critical situations
"""
import os
import asyncio
import aiohttp
from datetime import datetime, timedelta
import socketio
from fastapi import FastAPI, HTTPException, Depends, Header
import asyncpg
from contextlib import asynccontextmanager
# Configuration
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://adctrl:adctrl@postgres:5432/adctrl")
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"))
TELEGRAM_API_URL = os.getenv("TELEGRAM_API_URL", "http://tg-bot:8003/send")
# Database pool
db_pool = None
ws_task = None
# Auth dependency
async def verify_token(authorization: str = Header(None)):
if not authorization or not authorization.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Missing or invalid authorization header")
token = authorization.replace("Bearer ", "")
if token != SECRET_TOKEN:
raise HTTPException(status_code=403, detail="Invalid token")
return token
async def send_telegram_alert(message: str, service_id: int = None, service_name: str = None):
"""Send alert to telegram bot"""
try:
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=payload,
headers={"Authorization": f"Bearer {SECRET_TOKEN}"}
) as resp:
if resp.status != 200:
print(f"Failed to send telegram alert: Status {resp.status}")
except Exception as e:
print(f"Error sending telegram alert: {e}")
async def fetch_task_names():
"""Fetch task names from scoreboard API"""
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}
return {}
except Exception as e:
print(f"Error fetching task names: {e}")
return {}
async def socketio_listener():
"""Listen to ForcAD scoreboard using Socket.IO"""
sio = socketio.AsyncClient(logger=False, engineio_logger=False)
# Cache for task and team names
task_names = {}
team_names = {}
# Fetch task names on startup
task_names.update(await fetch_task_names())
@sio.on('*', namespace='/live_events')
async def catch_all(event, data):
"""Catch all events from live_events namespace"""
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) and 'data' in data:
await process_flag_stolen(data['data'])
async def process_flag_stolen(event_data):
"""Process flag_stolen event"""
try:
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)
if attacker_id is None or victim_id is None:
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"Flag event: attacker={attacker_id}, victim={victim_id}, service={service_name}, points={attacker_delta:.2f}")
print(f" Our team: {OUR_TEAM_ID}, is_our_attack={is_our_attack}, is_attack_to_us={is_attack_to_us}")
if is_our_attack or is_attack_to_us:
conn = await db_pool.acquire()
try:
attack_id = f"flag_{attacker_id}_{victim_id}_{task_id}_{int(timestamp.timestamp())}"
await conn.execute("""
INSERT INTO attacks (attack_id, attacker_team_id, victim_team_id, service_name, timestamp, points, is_our_attack, is_attack_to_us)
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_attack_to_us and attacker_delta >= ALERT_THRESHOLD_POINTS:
print(f" Sending alert: points {attacker_delta:.2f} >= threshold {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)
print(f" Alert sent successfully")
else:
if is_our_attack:
print(f" Our successful attack - no alert needed")
else:
print(f" Attack to us but below threshold: {attacker_delta:.2f} < {ALERT_THRESHOLD_POINTS}")
finally:
await db_pool.release(conn)
except Exception as e:
print(f"Error processing flag_stolen event: {e}")
@sio.event(namespace='/live_events')
async def update_scoreboard(data):
"""Handle scoreboard update - compare with previous state to detect NEW attacks"""
try:
event_data = data.get('data', {})
round_num = event_data.get('round', 0)
round_start = event_data.get('round_start', 0)
team_tasks = event_data.get('team_tasks', [])
conn = await db_pool.acquire()
try:
# Store team scores from team_tasks (score field = FP for this service)
# Aggregate scores per team
team_fp_totals = {}
for team_task in team_tasks:
team_id = team_task.get('team_id')
fp_score = team_task.get('score', 0)
if team_id not in team_fp_totals:
team_fp_totals[team_id] = 0
team_fp_totals[team_id] += fp_score
# Store aggregated scores
for team_id, total_fp in team_fp_totals.items():
await conn.execute("""
INSERT INTO team_scores (team_id, team_name, total_score, flag_points, round, timestamp)
VALUES ($1, $2, $3, $4, $5, NOW())
""", team_id, team_names.get(team_id, f'Team {team_id}'),
total_fp, total_fp, round_num)
# Process each team_task for attack detection
# Group by service to match stolen/lost pairs
service_data = {}
for team_task in team_tasks:
task_id = team_task.get('task_id')
service_name = task_names.get(task_id, f"task_{task_id}")
if service_name not in service_data:
service_data[service_name] = []
service_data[service_name].append(team_task)
# Process each service
for service_name, tasks in service_data.items():
# Track state for each team in this service
for team_task in tasks:
team_id = team_task.get('team_id')
task_id = team_task.get('task_id')
current_stolen = team_task.get('stolen', 0)
current_lost = team_task.get('lost', 0)
current_fp_score = team_task.get('score', 0)
# Get previous state from database
prev_state = await conn.fetchrow(
"SELECT stolen_flags, lost_flags, fp_score FROM scoreboard_state WHERE team_id = $1 AND service_name = $2",
team_id, service_name
)
prev_stolen = prev_state['stolen_flags'] if prev_state else 0
prev_lost = prev_state['lost_flags'] if prev_state else 0
prev_fp_score = prev_state['fp_score'] if prev_state else 0
# Calculate NEW flags and FP changes
new_stolen = current_stolen - prev_stolen
new_lost = current_lost - prev_lost
fp_change = current_fp_score - prev_fp_score
is_first_update = prev_state is None
# Update current state in database
await conn.execute("""
INSERT INTO scoreboard_state (team_id, service_name, stolen_flags, lost_flags, fp_score, last_updated)
VALUES ($1, $2, $3, $4, $5, NOW())
ON CONFLICT (team_id, service_name)
DO UPDATE SET stolen_flags = $3, lost_flags = $4, fp_score = $5, last_updated = NOW()
""", team_id, service_name, current_stolen, current_lost, current_fp_score)
# Create single attack record when flags change (not first update)
# Only track attacks involving our team to avoid duplicates
if not is_first_update and (new_stolen > 0 or new_lost > 0):
timestamp = datetime.utcnow()
is_our_attack = (new_stolen > 0 and team_id == OUR_TEAM_ID)
is_attack_to_us = (new_lost > 0 and team_id == OUR_TEAM_ID)
# Only create records for attacks involving OUR team
should_record = is_our_attack or is_attack_to_us
if should_record:
# Determine attacker/victim and FP
if new_stolen > 0:
# This team stole flags (attacker)
attacker_id = team_id
victim_id = None # We don't know exact victim
fp_value = max(0, fp_change)
attack_type = "stolen"
else:
# This team lost flags (victim)
attacker_id = None # We don't know exact attacker
victim_id = team_id
fp_value = abs(min(0, fp_change))
attack_type = "lost"
attack_id = f"r{round_num}_{attack_type}_team{team_id}_{service_name}_{int(timestamp.timestamp())}"
await conn.execute("""
INSERT INTO attacks (attack_id, attacker_team_id, victim_team_id, service_name, timestamp, points, is_our_attack, is_attack_to_us)
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(fp_value), is_our_attack, is_attack_to_us)
if is_our_attack:
pass
elif is_attack_to_us:
if fp_value >= ALERT_THRESHOLD_POINTS:
await check_and_create_alerts(conn, 0, service_name)
finally:
await db_pool.release(conn)
except Exception as e:
print(f"Error processing update_scoreboard: {e}")
@sio.event(namespace='/live_events')
async def init_scoreboard(data):
"""Handle initial scoreboard data"""
try:
event_data = data.get('data', {})
teams = event_data.get('teams', [])
tasks = event_data.get('tasks', [])
for task in tasks:
task_names[task.get('id')] = task.get('name')
for team in teams:
team_names[team.get('id')] = team.get('name')
except Exception as e:
print(f"Error processing init_scoreboard: {e}")
@sio.event
async def connect():
print(f"✅ Connected to scoreboard at {SCOREBOARD_URL}")
@sio.event
async def disconnect():
print(f"❌ Disconnected from scoreboard")
while True:
try:
print(f"Connecting to {SCOREBOARD_URL}...")
await sio.connect(
SCOREBOARD_URL,
namespaces=['/live_events'],
transports=['websocket']
)
await sio.wait()
except Exception as e:
print(f"Connection error: {e}")
await asyncio.sleep(5)
# Lifespan context
@asynccontextmanager
async def lifespan(app: FastAPI):
global db_pool, ws_task
db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
ws_task = asyncio.create_task(socketio_listener())
yield
if ws_task:
ws_task.cancel()
try:
await ws_task
except asyncio.CancelledError:
pass
await db_pool.close()
app = FastAPI(title="Scoreboard Injector", lifespan=lifespan)
# API Endpoints
@app.get("/health")
async def health_check():
return {
"status": "ok",
"timestamp": datetime.utcnow().isoformat(),
"team_id": OUR_TEAM_ID,
"mode": "socketio",
"scoreboard_url": SCOREBOARD_URL
}
@app.get("/stats", dependencies=[Depends(verify_token)])
async def get_stats():
"""Get attack statistics"""
conn = await get_db()
try:
total = await conn.fetchval("SELECT COUNT(*) FROM attacks")
attacks_by_us = await conn.fetchval("SELECT COUNT(*) FROM attacks WHERE is_our_attack = true")
attacks_to_us = await conn.fetchval("SELECT COUNT(*) FROM attacks WHERE is_attack_to_us = true")
threshold_time = datetime.utcnow() - timedelta(minutes=5)
recent = await conn.fetchval("SELECT COUNT(*) FROM attacks WHERE timestamp > $1", threshold_time)
critical_alerts = await conn.fetchval(
"SELECT COUNT(*) FROM attack_alerts WHERE severity = 'critical' AND created_at > $1",
threshold_time
)
return {
"total_attacks": total,
"attacks_by_us": attacks_by_us,
"attacks_to_us": attacks_to_us,
"recent_attacks_5min": recent,
"critical_alerts_5min": critical_alerts
}
finally:
await release_db(conn)
@app.get("/attacks", dependencies=[Depends(verify_token)])
async def get_attacks(limit: int = 100, our_attacks: Optional[bool] = None, attacks_to_us: Optional[bool] = None):
"""Get recent attacks with team names"""
conn = await get_db()
try:
query = """
SELECT
a.*,
ts_attacker.team_name as attacker_team_name,
ts_victim.team_name as victim_team_name
FROM attacks a
LEFT JOIN (
SELECT DISTINCT ON (team_id) team_id, team_name
FROM team_scores
ORDER BY team_id, timestamp DESC
) ts_attacker ON a.attacker_team_id = ts_attacker.team_id
LEFT JOIN (
SELECT DISTINCT ON (team_id) team_id, team_name
FROM team_scores
ORDER BY team_id, timestamp DESC
) ts_victim ON a.victim_team_id = ts_victim.team_id
WHERE 1=1
"""
params = []
param_count = 0
if our_attacks is not None:
param_count += 1
query += f" AND a.is_our_attack = ${param_count}"
params.append(our_attacks)
if attacks_to_us is not None:
param_count += 1
query += f" AND a.is_attack_to_us = ${param_count}"
params.append(attacks_to_us)
param_count += 1
query += f" ORDER BY a.timestamp DESC LIMIT ${param_count}"
params.append(limit)
rows = await conn.fetch(query, *params)
return [dict(row) for row in rows]
finally:
await release_db(conn)
@app.get("/alerts", dependencies=[Depends(verify_token)])
async def get_alerts(limit: int = 50, unnotified: bool = False):
"""Get alerts"""
conn = await get_db()
try:
if unnotified:
query = "SELECT * FROM attack_alerts WHERE notified = false ORDER BY created_at DESC LIMIT $1"
else:
query = "SELECT * FROM attack_alerts ORDER BY created_at DESC LIMIT $1"
rows = await conn.fetch(query, limit)
return [dict(row) for row in rows]
finally:
await release_db(conn)
@app.post("/alerts/{alert_id}/acknowledge", dependencies=[Depends(verify_token)])
async def acknowledge_alert(alert_id: int):
"""Mark alert as acknowledged"""
conn = await get_db()
try:
await conn.execute("UPDATE attack_alerts SET notified = true WHERE id = $1", alert_id)
return {"status": "acknowledged", "alert_id": alert_id}
finally:
await release_db(conn)
@app.get("/attacks/by-service", dependencies=[Depends(verify_token)])
async def get_attacks_by_service():
"""Get attack statistics grouped by service"""
conn = await get_db()
try:
rows = await conn.fetch("""
SELECT
service_name,
COUNT(*) as total_attacks,
COUNT(*) FILTER (WHERE is_our_attack = true) as our_attacks,
COUNT(*) FILTER (WHERE is_attack_to_us = true) as attacks_to_us,
COALESCE(SUM(points) FILTER (WHERE is_our_attack = true), 0) as points_gained,
COALESCE(SUM(points) FILTER (WHERE is_attack_to_us = true), 0) as points_lost
FROM attacks
GROUP BY service_name
ORDER BY total_attacks DESC
""")
return [dict(row) for row in rows]
finally:
await release_db(conn)
@app.post("/settings/team-id", dependencies=[Depends(verify_token)])
async def set_team_id(team_id: int):
"""Update our team ID"""
global OUR_TEAM_ID
OUR_TEAM_ID = team_id
conn = await get_db()
try:
await conn.execute(
"INSERT INTO settings (key, value) VALUES ('our_team_id', $1) ON CONFLICT (key) DO UPDATE SET value = $1",
str(team_id)
)
return {"team_id": team_id}
finally:
await release_db(conn)
@app.get("/settings/team-id", dependencies=[Depends(verify_token)])
async def get_team_id():
"""Get current team ID setting"""
return {"team_id": OUR_TEAM_ID}
@app.post("/test/inject-attack", dependencies=[Depends(verify_token)])
async def inject_test_attack(attacker_id: int, victim_id: int, service: str = "test-service", points: float = 10.0):
"""Manually inject a test attack event for debugging"""
test_event = {
"type": "attack",
"attacker_id": attacker_id,
"victim_id": victim_id,
"service": service,
"flag": "TEST_FLAG_" + datetime.utcnow().isoformat(),
"points": points,
"time": datetime.utcnow().isoformat(),
"round": 1
}
await process_attack_event(test_event)
return {"status": "injected", "event": test_event}
@app.get("/debug/scoreboard", dependencies=[Depends(verify_token)])
async def debug_scoreboard():
"""Check if scoreboard is reachable and show connection info"""
import aiohttp
results = {
"mode": "socketio",
"config": {
"scoreboard_url": SCOREBOARD_URL,
"our_team_id": OUR_TEAM_ID
},
"endpoints_tested": []
}
try:
async with aiohttp.ClientSession() as session:
# Test Socket.IO endpoint
socketio_url = f"{SCOREBOARD_URL}/socket.io/?EIO=4&transport=polling"
try:
async with session.get(socketio_url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
results["socketio_status"] = {
"url": socketio_url,
"status": resp.status,
"reachable": resp.status == 200,
"response_preview": (await resp.text())[:200] if resp.status == 200 else None
}
except Exception as e:
results["socketio_status"] = {
"url": socketio_url,
"reachable": False,
"error": str(e)
}
# Test base scoreboard URL
try:
async with session.get(SCOREBOARD_URL, timeout=aiohttp.ClientTimeout(total=5)) as resp:
results["base_url_status"] = {
"url": SCOREBOARD_URL,
"status": resp.status,
"reachable": resp.status == 200
}
except Exception as e:
results["base_url_status"] = {
"url": SCOREBOARD_URL,
"reachable": False,
"error": str(e)
}
# Test attack_data endpoint (for reference only)
attack_data_url = f"{SCOREBOARD_URL}/api/client/attack_data"
try:
async with session.get(attack_data_url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
result = {
"url": attack_data_url,
"status": resp.status,
"reachable": resp.status == 200,
"content_type": resp.headers.get('Content-Type', ''),
"note": "Contains exploit credentials, not attack events"
}
if resp.status == 200 and 'application/json' in resp.headers.get('Content-Type', ''):
data = await resp.json()
result["services"] = list(data.keys()) if isinstance(data, dict) else None
results["endpoints_tested"].append(result)
except Exception as e:
results["endpoints_tested"].append({
"url": attack_data_url,
"reachable": False,
"error": str(e)
})
except Exception as e:
results["error"] = str(e)
return results
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8002)