Update main.py

This commit is contained in:
ilyastar9999
2025-12-03 10:22:19 +03:00
parent 462363e52f
commit 49ccad4e37

View File

@@ -55,29 +55,74 @@ async def process_attack_event(event: Dict[str, Any]):
conn = await db_pool.acquire() conn = await db_pool.acquire()
try: try:
# Extract attack information from event # Extract attack information from event
# Adjust fields based on actual ForcAD event structure # Handle multiple possible event formats from ForcAD
attack_id = event.get('id') or f"{event.get('round')}_{event.get('attacker_id')}_{event.get('victim_id')}_{event.get('service')}" event_type = event.get('type', 'unknown')
attacker_id = event.get('attacker_id') or event.get('team_id')
victim_id = event.get('victim_id') or event.get('target_id') # Try to extract attacker/victim IDs from various possible fields
service_name = event.get('service') or event.get('service_name') 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', '') flag = event.get('flag', '')
timestamp = datetime.fromisoformat(event.get('time', datetime.utcnow().isoformat()))
points = float(event.get('points', 0)) # 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_our_attack = attacker_id == OUR_TEAM_ID
is_attack_to_us = victim_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 # Store attack in database
await conn.execute(""" 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) 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) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
ON CONFLICT (attack_id) DO NOTHING 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) """, 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 # Check for alert conditions if attack is against us
if is_attack_to_us: if is_attack_to_us:
await check_and_create_alerts(conn, attacker_id, service_name) 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: finally:
await db_pool.release(conn) await db_pool.release(conn)
@@ -141,30 +186,61 @@ async def send_telegram_alert(message: str):
async def websocket_listener(): async def websocket_listener():
"""Listen to scoreboard WebSocket for events""" """Listen to scoreboard WebSocket for events"""
reconnect_delay = 5
max_reconnect_delay = 60
while True: while True:
try: try:
print(f"Attempting to connect to scoreboard WebSocket: {SCOREBOARD_WS_URL}")
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.ws_connect(SCOREBOARD_WS_URL) as ws: try:
print(f"Connected to scoreboard WebSocket: {SCOREBOARD_WS_URL}") async with session.ws_connect(
SCOREBOARD_WS_URL,
timeout=aiohttp.ClientTimeout(total=10)
) as ws:
print(f"✓ Connected to scoreboard WebSocket")
reconnect_delay = 5 # Reset delay on successful connection
async for msg in ws: async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT: if msg.type == aiohttp.WSMsgType.TEXT:
try: try:
event = json.loads(msg.data) event = json.loads(msg.data)
# Process different event types event_type = event.get('type', 'unknown')
if event.get('type') in ['attack', 'flag_stolen', 'service_status']:
# Process attack-related events
# ForcAD can send: 'attack', 'flag_stolen', 'stolen_flag', 'flag_submit'
if event_type in ['attack', 'flag_stolen', 'stolen_flag', 'flag_submit', 'flag']:
await process_attack_event(event) await process_attack_event(event)
except json.JSONDecodeError: # Optionally log other event types for debugging
print(f"Failed to decode WebSocket message: {msg.data}") # else:
# print(f"Received event type: {event_type}")
except json.JSONDecodeError as e:
print(f"Failed to decode WebSocket message: {msg.data[:200]}")
except Exception as e: except Exception as e:
print(f"Error processing event: {e}") print(f"Error processing event: {e}")
elif msg.type == aiohttp.WSMsgType.CLOSED:
print("WebSocket connection closed by server")
break
elif msg.type == aiohttp.WSMsgType.ERROR: elif msg.type == aiohttp.WSMsgType.ERROR:
print(f"WebSocket error: {ws.exception()}") print(f"WebSocket error: {ws.exception()}")
break break
except aiohttp.ClientError as e:
print(f"WebSocket client error: {e}")
except asyncio.TimeoutError:
print(f"Connection timeout to {SCOREBOARD_WS_URL}")
except Exception as e: except Exception as e:
print(f"WebSocket connection error: {e}") print(f"WebSocket connection error: {e}")
print("Reconnecting in 5 seconds...")
await asyncio.sleep(5) print(f"Reconnecting in {reconnect_delay} seconds...")
await asyncio.sleep(reconnect_delay)
# Exponential backoff
reconnect_delay = min(reconnect_delay * 1.5, max_reconnect_delay)
# Lifespan context # Lifespan context
@asynccontextmanager @asynccontextmanager
@@ -311,6 +387,28 @@ async def set_team_id(team_id: int):
finally: finally:
await release_db(conn) 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}
if __name__ == "__main__": if __name__ == "__main__":
import uvicorn import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8002) uvicorn.run(app, host="0.0.0.0", port=8002)