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

11
controler/main.py
View File

@@ -3,12 +3,10 @@ API Controller for A/D Infrastructure
Manages docker-compose services with authentication Manages docker-compose services with authentication
""" """
import os import os
import subprocess
import asyncio import asyncio
from datetime import datetime from datetime import datetime
from typing import Optional, List from typing import Optional, List, Tuple
from fastapi import FastAPI, HTTPException, Depends, Header from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel from pydantic import BaseModel
import asyncpg import asyncpg
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
@@ -27,14 +25,11 @@ class ServiceCreate(BaseModel):
git_url: Optional[str] = None git_url: Optional[str] = None
class ServiceAction(BaseModel): class ServiceAction(BaseModel):
action: str # start, stop, restart action: str
class GitPullRequest(BaseModel): class GitPullRequest(BaseModel):
auto: bool = False auto: bool = False
class LogRequest(BaseModel):
lines: int = 100
# Auth dependency # Auth dependency
async def verify_token(authorization: str = Header(None)): async def verify_token(authorization: str = Header(None)):
if not authorization or not authorization.startswith("Bearer "): if not authorization or not authorization.startswith("Bearer "):
@@ -114,7 +109,7 @@ def parse_docker_status(ps_output: str) -> dict:
'containers': containers '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)""" """Run docker-compose command and return (returncode, stdout, stderr)"""
compose_file = find_compose_file(service_path) compose_file = find_compose_file(service_path)
if not compose_file: if not compose_file:

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 Monitors Socket.IO events for attacks and alerts on critical situations
""" """
import os import os
import json
import asyncio import asyncio
import aiohttp
from datetime import datetime, timedelta from datetime import datetime, timedelta
from typing import Optional, Dict, Any
import socketio import socketio
from fastapi import FastAPI, HTTPException, Depends, Header from fastapi import FastAPI, HTTPException, Depends, Header
from pydantic import BaseModel
import asyncpg import asyncpg
from contextlib import asynccontextmanager 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") SCOREBOARD_URL = os.getenv("SCOREBOARD_URL", "http://10.60.0.1:8080")
OUR_TEAM_ID = int(os.getenv("OUR_TEAM_ID", "1")) OUR_TEAM_ID = int(os.getenv("OUR_TEAM_ID", "1"))
ALERT_THRESHOLD_POINTS = float(os.getenv("ALERT_THRESHOLD_POINTS", "5")) 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") TELEGRAM_API_URL = os.getenv("TELEGRAM_API_URL", "http://tg-bot:8003/send")
# Database pool # Database pool
db_pool = None db_pool = None
ws_task = 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 # Auth dependency
async def verify_token(authorization: str = Header(None)): async def verify_token(authorization: str = Header(None)):
if not authorization or not authorization.startswith("Bearer "): 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") raise HTTPException(status_code=403, detail="Invalid token")
return token return token
# Database functions
async def get_db():
return await db_pool.acquire()
async def release_db(conn): async def send_telegram_alert(message: str, service_id: int = None, service_name: str = None):
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):
"""Send alert to telegram bot""" """Send alert to telegram bot"""
import aiohttp
try: try:
print(f"📱 Sending alert to Telegram: {TELEGRAM_API_URL}")
async with aiohttp.ClientSession() as session: 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( async with session.post(
TELEGRAM_API_URL, TELEGRAM_API_URL,
json={"message": message}, json=payload,
headers={"Authorization": f"Bearer {SECRET_TOKEN}"} headers={"Authorization": f"Bearer {SECRET_TOKEN}"}
) as resp: ) as resp:
response_text = await resp.text()
if resp.status != 200: if resp.status != 200:
print(f"❌ Failed to send telegram alert: Status {resp.status}") 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}")
except Exception as e: except Exception as e:
print(f"❌ Error sending telegram alert: {e}") print(f"Error sending telegram alert: {e}")
import traceback
traceback.print_exc()
async def fetch_task_names(): async def fetch_task_names():
"""Fetch task names from scoreboard API""" """Fetch task names from scoreboard API"""
import aiohttp
try: try:
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.get(f"{SCOREBOARD_URL}/api/client/tasks/") as resp: async with session.get(f"{SCOREBOARD_URL}/api/client/tasks/") as resp:
if resp.status == 200: if resp.status == 200:
tasks = await resp.json() tasks = await resp.json()
return {task['id']: task['name'] for task in tasks} return {task['id']: task['name'] for task in tasks}
else: return {}
print(f"Failed to fetch tasks: {resp.status}")
return {}
except Exception as e: except Exception as e:
print(f"Error fetching task names: {e}") print(f"Error fetching task names: {e}")
return {} return {}
@@ -219,44 +77,35 @@ async def socketio_listener():
# Fetch task names on startup # Fetch task names on startup
task_names.update(await fetch_task_names()) 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') @sio.on('*', namespace='/live_events')
async def catch_all(event, data): async def catch_all(event, data):
"""Catch all events from live_events namespace""" """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: if isinstance(data, list) and len(data) >= 2:
event_type = data[0] event_type = data[0]
event_data = data[1].get('data', {}) if isinstance(data[1], dict) else {} event_data = data[1].get('data', {}) if isinstance(data[1], dict) else {}
if event_type == 'flag_stolen': if event_type == 'flag_stolen':
await process_flag_stolen(event_data) await process_flag_stolen(event_data)
elif isinstance(data, dict): elif isinstance(data, dict) and 'data' in data:
# Handle direct event data await process_flag_stolen(data['data'])
if 'data' in data:
await process_flag_stolen(data['data'])
async def process_flag_stolen(event_data): async def process_flag_stolen(event_data):
"""Process flag_stolen event""" """Process flag_stolen event"""
try: try:
print(f"[DEBUG] process_flag_stolen called with event_data: {event_data}")
attacker_id = event_data.get('attacker_id') attacker_id = event_data.get('attacker_id')
victim_id = event_data.get('victim_id') victim_id = event_data.get('victim_id')
task_id = event_data.get('task_id') task_id = event_data.get('task_id')
attacker_delta = event_data.get('attacker_delta', 0) 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: if attacker_id is None or victim_id is None:
print("[DEBUG] attacker_id or victim_id is None, skipping event")
return return
service_name = task_names.get(task_id, f"task_{task_id}") service_name = task_names.get(task_id, f"task_{task_id}")
timestamp = datetime.utcnow() timestamp = datetime.utcnow()
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
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: if is_our_attack or is_attack_to_us:
conn = await db_pool.acquire() conn = await db_pool.acquire()
try: try:
@@ -266,34 +115,38 @@ async def socketio_listener():
VALUES ($1, $2, $3, $4, $5, $6, $7, $8) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
ON CONFLICT (attack_id) DO NOTHING 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) """, 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)") if is_attack_to_us and attacker_delta >= ALERT_THRESHOLD_POINTS:
elif is_attack_to_us: alert_message = f"🚨 ATTACK DETECTED!\nTeam {attacker_id} stole flag from {service_name}\nPoints lost: {attacker_delta:.2f} FP"
print(f" ⚠️ Team {attacker_id} stole flag from us on {service_name} (-{attacker_delta:.2f} FP)")
if attacker_delta >= ALERT_THRESHOLD_POINTS: # Get service_id from controller if available
print(f"[DEBUG] Sending Telegram alert: attacker_delta={attacker_delta} >= ALERT_THRESHOLD_POINTS={ALERT_THRESHOLD_POINTS}") service_id = None
alert_message = f"🚨 ATTACK DETECTED!\nTeam {attacker_id} stole flag from {service_name}\nPoints lost: {attacker_delta:.2f} FP" try:
alert_id = await conn.fetchval(""" service_row = await conn.fetchrow(
INSERT INTO attack_alerts (attack_id, alert_type, severity, message) "SELECT id FROM services WHERE name = $1 LIMIT 1",
VALUES ( service_name
(SELECT id FROM attacks WHERE attack_id = $1), )
'flag_stolen', if service_row:
'high', service_id = service_row['id']
$2 except:
) pass
RETURNING id
""", attack_id, alert_message) alert_id = await conn.fetchval("""
await send_telegram_alert(alert_message) INSERT INTO attack_alerts (attack_id, alert_type, severity, message)
await conn.execute("UPDATE attack_alerts SET notified = true WHERE id = $1", alert_id) VALUES (
print(f" 📱 Alert sent to Telegram") (SELECT id FROM attacks WHERE attack_id = $1),
else: 'flag_stolen',
print(f"[DEBUG] No alert sent: attacker_delta={attacker_delta} < ALERT_THRESHOLD_POINTS={ALERT_THRESHOLD_POINTS}") '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: finally:
await db_pool.release(conn) await db_pool.release(conn)
except Exception as e: except Exception as e:
print(f"Error processing flag_stolen event: {e}") print(f"Error processing flag_stolen event: {e}")
import traceback
traceback.print_exc()
@sio.event(namespace='/live_events') @sio.event(namespace='/live_events')
async def update_scoreboard(data): async def update_scoreboard(data):
@@ -304,8 +157,6 @@ async def socketio_listener():
round_start = event_data.get('round_start', 0) round_start = event_data.get('round_start', 0)
team_tasks = event_data.get('team_tasks', []) team_tasks = event_data.get('team_tasks', [])
print(f"📊 Round {round_num} - Processing {len(team_tasks)} team updates")
conn = await db_pool.acquire() conn = await db_pool.acquire()
try: try:
# Store team scores from team_tasks (score field = FP for this service) # 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) """, attack_id, attacker_id, victim_id, service_name, timestamp, float(fp_value), is_our_attack, is_attack_to_us)
if is_our_attack: if is_our_attack:
print(f" ✅ We stole {new_stolen} flags from {service_name} (+{fp_value:.2f} FP)") pass
elif is_attack_to_us: 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: if fp_value >= ALERT_THRESHOLD_POINTS:
await check_and_create_alerts(conn, 0, service_name) await check_and_create_alerts(conn, 0, service_name)
elif new_stolen > 0: finally:
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:
await db_pool.release(conn) await db_pool.release(conn)
except Exception as e: except Exception as e:
print(f"Error processing update_scoreboard: {e}") print(f"Error processing update_scoreboard: {e}")
import traceback
traceback.print_exc()
@sio.event(namespace='/live_events') @sio.event(namespace='/live_events')
async def init_scoreboard(data): async def init_scoreboard(data):
"""Handle initial scoreboard data""" """Handle initial scoreboard data"""
try: try:
print("📡 Received initial scoreboard data")
event_data = data.get('data', {}) event_data = data.get('data', {})
teams = event_data.get('teams', []) teams = event_data.get('teams', [])
tasks = event_data.get('tasks', []) tasks = event_data.get('tasks', [])
# Cache task names
for task in tasks: for task in tasks:
task_names[task.get('id')] = task.get('name') task_names[task.get('id')] = task.get('name')
# Cache team names
for team in teams: for team in teams:
team_names[team.get('id')] = team.get('name') 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: except Exception as e:
print(f"Error processing init_scoreboard: {e}") print(f"Error processing init_scoreboard: {e}")
@sio.event @sio.event
async def connect(): async def connect():
print("✅ Connected to ForcAD scoreboard Socket.IO") pass
@sio.event @sio.event
async def disconnect(): async def disconnect():
print("❌ Disconnected from scoreboard") pass
while True: while True:
try: try:
print(f"Connecting to {SCOREBOARD_URL}/socket.io ...")
await sio.connect( await sio.connect(
SCOREBOARD_URL, SCOREBOARD_URL,
namespaces=['/live_events'], namespaces=['/live_events'],
@@ -468,8 +303,6 @@ async def socketio_listener():
) )
await sio.wait() await sio.wait()
except Exception as e: except Exception as e:
print(f"Socket.IO error: {e}")
print("Reconnecting in 5 seconds...")
await asyncio.sleep(5) await asyncio.sleep(5)
# Lifespan context # Lifespan context
@@ -477,16 +310,10 @@ async def socketio_listener():
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
global db_pool, ws_task global db_pool, ws_task
db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10) 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()) ws_task = asyncio.create_task(socketio_listener())
yield yield
# Cleanup
if ws_task: if ws_task:
ws_task.cancel() ws_task.cancel()
try: try:

105
tg-bot/main.py
View File

@@ -7,8 +7,9 @@ from datetime import datetime
from fastapi import FastAPI, HTTPException, Depends, Header from fastapi import FastAPI, HTTPException, Depends, Header
from pydantic import BaseModel from pydantic import BaseModel
import asyncpg import asyncpg
from telegram import Bot from telegram import Bot, InlineKeyboardButton, InlineKeyboardMarkup, Update
from telegram.error import TelegramError from telegram.error import TelegramError
from telegram.ext import Application, CallbackQueryHandler, ContextTypes
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
# Configuration # Configuration
@@ -17,19 +18,56 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production")
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "") TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID", "") 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 # Database pool and bot
db_pool = None db_pool = None
bot = 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): class MessageRequest(BaseModel):
message: str 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): class BulkMessageRequest(BaseModel):
messages: list[str] messages: list[str]
@@ -66,14 +104,22 @@ async def log_message(chat_id: int, message: str, success: bool, error_message:
# Lifespan context # Lifespan context
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): 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) db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
if TELEGRAM_BOT_TOKEN: if TELEGRAM_BOT_TOKEN:
bot = Bot(token=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 yield
if app_telegram:
await app_telegram.shutdown()
await db_pool.close() await db_pool.close()
app = FastAPI(title="Telegram Bot API", lifespan=lifespan) app = FastAPI(title="Telegram Bot API", lifespan=lifespan)
@@ -89,7 +135,7 @@ async def health_check():
@app.post("/send", dependencies=[Depends(verify_token)]) @app.post("/send", dependencies=[Depends(verify_token)])
async def send_message(request: MessageRequest): 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: if not bot:
raise HTTPException(status_code=503, detail="Telegram bot not configured") 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") raise HTTPException(status_code=400, detail="No chat_id provided and no default configured")
try: try:
# Send message kwargs = {
message = await bot.send_message( "chat_id": int(chat_id),
chat_id=int(chat_id), "text": request.message,
text=request.message, "parse_mode": "HTML"
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) await log_message(int(chat_id), request.message, True)
return { return {
@@ -115,7 +171,6 @@ async def send_message(request: MessageRequest):
} }
except TelegramError as e: except TelegramError as e:
# Log failure
await log_message(int(chat_id), request.message, False, str(e)) await log_message(int(chat_id), request.message, False, str(e))
raise HTTPException(status_code=500, detail=f"Failed to send message: {str(e)}") raise HTTPException(status_code=500, detail=f"Failed to send message: {str(e)}")
@@ -184,6 +239,20 @@ async def get_stats():
finally: finally:
await release_db(conn) 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)]) @app.post("/test", dependencies=[Depends(verify_token)])
async def test_connection(): async def test_connection():
"""Test telegram bot connection""" """Test telegram bot connection"""

3
tg-bot/requirements.txt
View File

@@ -2,5 +2,6 @@ fastapi==0.109.0
uvicorn[standard]==0.27.0 uvicorn[standard]==0.27.0
asyncpg==0.29.0 asyncpg==0.29.0
pydantic==2.5.3 pydantic==2.5.3
python-telegram-bot==21.0 python-telegram-bot[all]==21.0
python-dotenv==1.0.0 python-dotenv==1.0.0
aiohttp==3.9.1

8
web/app.py
View File

@@ -4,14 +4,12 @@ Flask-based dashboard to monitor services, attacks, and alerts
""" """
import os import os
import asyncio import asyncio
from datetime import datetime, timedelta from datetime import datetime
from flask import Flask, render_template, jsonify, request, redirect, url_for, session from flask import Flask, render_template, jsonify, request, redirect, url_for, session
import asyncpg
import aiohttp import aiohttp
from functools import wraps from functools import wraps
# Configuration # Configuration
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://adctrl:adctrl@postgres:5432/adctrl")
SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production") SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production")
WEB_PASSWORD = os.getenv("WEB_PASSWORD", "admin123") WEB_PASSWORD = os.getenv("WEB_PASSWORD", "admin123")
CONTROLLER_API = os.getenv("CONTROLLER_API", "http://controller:8001") 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 = Flask(__name__)
app.secret_key = os.getenv("FLASK_SECRET_KEY", "change-me-in-production-flask-secret") 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 # Auth decorator
def login_required(f): def login_required(f):
@wraps(f) @wraps(f)