Botnet deployments work!

:yayayayayay-ctp:

1 files changed, 99 insertions, 0 deletions

diff --git a/src/db.py b/src/db.py
index 6dd5cad..9e1619a 100644
--- a/src/db.py
+++ b/src/db.py
@@ -1,4 +1,5 @@
 from slack_sdk.web.async_client import AsyncWebClient
+from datetime import datetime
 import aiosqlite
 
 DB_PATH = "botnet.db"
@@ -12,6 +13,28 @@ async def init_db():
                 joined_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
             )
         """)
+        await db.execute("""
+             CREATE TABLE IF NOT EXISTS Deployments (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                message TEXT NOT NULL,
+                submitted_by TEXT NOT NULL,
+                status TEXT DEFAULT 'pending',
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                deployed_at TIMESTAMP
+             )
+         """)
+        await db.execute("""
+             CREATE TABLE IF NOT EXISTS DeploymentMessages (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                deployment_id INTEGER NOT NULL,
+                user_id TEXT NOT NULL,
+                message_id TEXT NOT NULL,
+                channel_id TEXT NOT NULL,
+                sent_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (deployment_id) REFERENCES Deployments(id) ON DELETE CASCADE,
+                FOREIGN KEY (user_id) REFERENCES Users(slack_id)
+             )
+         """)
         await db.commit()
 
 async def add_user(slack_id: str, oauth_token: str):
@@ -30,3 +53,79 @@ async def add_user(slack_id: str, oauth_token: str):
         )
     except Exception as e:
         print(f"Failed to send welcome message: {e}")
+
+async def create_deployment(user_id: str, message: str) -> int:
+    async with aiosqlite.connect(DB_PATH) as db:
+        cursor = await db.execute(
+            "INSERT INTO Deployments (message, submitted_by) VALUES (?, ?)",
+            (message, user_id)
+        )
+        await db.commit()
+        return cursor.lastrowid
+
+async def get_daily_user_deployment_count(user_id: str) -> int:
+    from datetime import datetime
+    today = datetime.now().date().isoformat()
+    async with aiosqlite.connect(DB_PATH) as db:
+        cursor = await db.execute(
+            """SELECT COUNT(*) FROM Deployments 
+               WHERE submitted_by = ? AND date(created_at) = ?""",
+            (user_id, today)
+        )
+        row = await cursor.fetchone()
+        return row[0] if row else 0
+
+async def get_deployment(deployment_id: int) -> dict | None:
+    async with aiosqlite.connect(DB_PATH) as db:
+        db.row_factory = aiosqlite.Row
+        cursor = await db.execute(
+            "SELECT * FROM Deployments WHERE id = ?",
+            (deployment_id,)
+        )
+        row = await cursor.fetchone()
+        return dict(row) if row else None
+
+async def update_deployment_status(deployment_id: int, status: str):
+    async with aiosqlite.connect(DB_PATH) as db:
+        await db.execute(
+            "UPDATE Deployments SET status = ? WHERE id = ?",
+            (status, deployment_id)
+        )
+        await db.commit()
+
+async def add_deployment_message(deployment_id: int, user_id: str, channel_id: str, message_id: str):
+    async with aiosqlite.connect(DB_PATH) as db:
+        await db.execute(
+            """INSERT INTO DeploymentMessages 
+               (deployment_id, user_id, channel_id, message_id) 
+               VALUES (?, ?, ?, ?)""",
+            (deployment_id, user_id, channel_id, message_id)
+        )
+        await db.commit()
+
+async def get_deployment_messages(deployment_id: int) -> list[dict]:
+    async with aiosqlite.connect(DB_PATH) as db:
+        db.row_factory = aiosqlite.Row
+        cursor = await db.execute(
+            "SELECT * FROM DeploymentMessages WHERE deployment_id = ?",
+            (deployment_id,)
+        )
+        rows = await cursor.fetchall()
+        return [dict(row) for row in rows]
+
+async def get_user(slack_id: str) -> dict | None:
+    async with aiosqlite.connect(DB_PATH) as db:
+        db.row_factory = aiosqlite.Row
+        cursor = await db.execute(
+            "SELECT * FROM Users WHERE slack_id = ?",
+            (slack_id,)
+        )
+        row = await cursor.fetchone()
+        return dict(row) if row else None
+
+async def get_active_users() -> list[dict]:
+    async with aiosqlite.connect(DB_PATH) as db:
+        db.row_factory = aiosqlite.Row
+        cursor = await db.execute("SELECT * FROM Users")
+        rows = await cursor.fetchall()
+        return [dict(row) for row in rows]