aboutsummaryrefslogtreecommitdiff
path: root/src/db.py
blob: 9e1619a45e3b9a3196da070c3dfefdfa17a24caf (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
from slack_sdk.web.async_client import AsyncWebClient
from datetime import datetime
import aiosqlite

DB_PATH = "botnet.db"

async def init_db():
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS Users (
                slack_id TEXT PRIMARY KEY,
                oauth_token TEXT NOT NULL,
                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):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR REPLACE INTO Users (slack_id, oauth_token) VALUES (?, ?)",
            (slack_id, oauth_token)
        )
        await db.commit()
    client = AsyncWebClient(token=oauth_token)
    try:
        await client.chat_postMessage(
            channel="C0BGMV1CAQG",
            text="I've joined the botnet!",
            as_user=True
        )
    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]