135 lines
5.2 KiB
Python
135 lines
5.2 KiB
Python
import aiosqlite
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import List, Dict, Any, Optional
|
|
|
|
DB_PATH = "data/context_hub.db"
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
async def init_db():
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS rules (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
scope TEXT UNIQUE NOT NULL,
|
|
content TEXT NOT NULL
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS api_keys (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
agent_name TEXT UNIQUE NOT NULL,
|
|
api_key TEXT UNIQUE NOT NULL,
|
|
last_used TIMESTAMP
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS audit_log (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
agent_name TEXT,
|
|
scope TEXT,
|
|
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
ip_address TEXT
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS ports (
|
|
port INTEGER PRIMARY KEY,
|
|
service_name TEXT NOT NULL,
|
|
description TEXT
|
|
)
|
|
""")
|
|
await db.commit()
|
|
|
|
async def get_rule(scope: str) -> Optional[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
async with db.execute("SELECT content FROM rules WHERE scope = ?", (scope,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
if row:
|
|
return json.loads(row[0])
|
|
return None
|
|
|
|
async def set_rule(scope: str, content: Dict[str, Any]):
|
|
content_str = json.dumps(content)
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO rules (scope, content) VALUES (?, ?) ON CONFLICT(scope) DO UPDATE SET content=?",
|
|
(scope, content_str, content_str)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_all_rules() -> Dict[str, Any]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
async with db.execute("SELECT scope, content FROM rules") as cursor:
|
|
rows = await cursor.fetchall()
|
|
return {row[0]: json.loads(row[1]) for row in rows}
|
|
|
|
async def get_all_ports() -> List[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT port, service_name, description FROM ports") as cursor:
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
async def add_port(port: int, service_name: str, description: str = ""):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO ports (port, service_name, description) VALUES (?, ?, ?) ON CONFLICT(port) DO UPDATE SET service_name=?, description=?",
|
|
(port, service_name, description, service_name, description)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_port(port: int) -> Optional[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT port, service_name, description FROM ports WHERE port = ?", (port,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
if row:
|
|
return dict(row)
|
|
return None
|
|
|
|
async def add_api_key(agent_name: str, api_key: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO api_keys (agent_name, api_key) VALUES (?, ?) ON CONFLICT(agent_name) DO UPDATE SET api_key=?",
|
|
(agent_name, api_key, api_key)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_agent_by_key(api_key: str) -> Optional[str]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
async with db.execute("SELECT agent_name FROM api_keys WHERE api_key = ?", (api_key,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
if row:
|
|
return row[0]
|
|
return None
|
|
|
|
async def update_key_last_used(api_key: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute("UPDATE api_keys SET last_used = CURRENT_TIMESTAMP WHERE api_key = ?", (api_key,))
|
|
await db.commit()
|
|
|
|
async def get_all_keys() -> List[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT id, agent_name, api_key, last_used FROM api_keys") as cursor:
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
async def add_audit_log(agent_name: str, scope: str, ip_address: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO audit_log (agent_name, scope, ip_address) VALUES (?, ?, ?)",
|
|
(agent_name, scope, ip_address)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_audit_logs(limit: int = 50) -> List[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT id, agent_name, scope, timestamp, ip_address FROM audit_log ORDER BY timestamp DESC LIMIT ?", (limit,)) as cursor:
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|