Update main.py
This commit is contained in:
ilyastar9999
@@ -55,29 +55,74 @@ async def process_attack_event(event: Dict[str, Any]):
|
||||
conn = await db_pool.acquire()
|
||||
try:
|
||||
# Extract attack information from event
|
||||
# Adjust fields based on actual ForcAD event structure
|
||||
attack_id = event.get('id') or f"{event.get('round')}_{event.get('attacker_id')}_{event.get('victim_id')}_{event.get('service')}"
|
||||
attacker_id = event.get('attacker_id') or event.get('team_id')
|
||||
victim_id = event.get('victim_id') or event.get('target_id')
|
||||
service_name = event.get('service') or event.get('service_name')
|
||||
# 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', '')
|
||||
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_attack_to_us = victim_id == OUR_TEAM_ID
|
||||
|
||||
# Store attack in database
|
||||
await conn.execute("""
|
||||
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
|
||||
""", attack_id, attacker_id, victim_id, service_name, flag, timestamp, points, is_our_attack, is_attack_to_us)
|
||||
|
||||
# Check for alert conditions if attack is against us
|
||||
if is_attack_to_us:
|
||||
await check_and_create_alerts(conn, attacker_id, service_name)
|
||||
# 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)
|
||||
|
||||
@@ -141,30 +186,61 @@ async def send_telegram_alert(message: str):
|
||||
|
||||
async def websocket_listener():
|
||||
"""Listen to scoreboard WebSocket for events"""
|
||||
reconnect_delay = 5
|
||||
max_reconnect_delay = 60
|
||||
|
||||
while True:
|
||||
try:
|
||||
print(f"Attempting to connect to scoreboard WebSocket: {SCOREBOARD_WS_URL}")
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.ws_connect(SCOREBOARD_WS_URL) as ws:
|
||||
print(f"Connected to scoreboard WebSocket: {SCOREBOARD_WS_URL}")
|
||||
try:
|
||||
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:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
event = json.loads(msg.data)
|
||||
event_type = event.get('type', 'unknown')
|
||||
|
||||
# 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)
|
||||
# Optionally log other event types for debugging
|
||||
# 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:
|
||||
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:
|
||||
print(f"WebSocket error: {ws.exception()}")
|
||||
break
|
||||
|
||||
except aiohttp.ClientError as e:
|
||||
print(f"WebSocket client error: {e}")
|
||||
except asyncio.TimeoutError:
|
||||
print(f"Connection timeout to {SCOREBOARD_WS_URL}")
|
||||
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
event = json.loads(msg.data)
|
||||
# Process different event types
|
||||
if event.get('type') in ['attack', 'flag_stolen', 'service_status']:
|
||||
await process_attack_event(event)
|
||||
except json.JSONDecodeError:
|
||||
print(f"Failed to decode WebSocket message: {msg.data}")
|
||||
except Exception as e:
|
||||
print(f"Error processing event: {e}")
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
print(f"WebSocket error: {ws.exception()}")
|
||||
break
|
||||
except Exception as 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
|
||||
@asynccontextmanager
|
||||
@@ -311,6 +387,28 @@ async def set_team_id(team_id: int):
|
||||
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}
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8002)
|
||||
|
||||
Reference in New Issue
Block a user