aboutsummaryrefslogtreecommitdiff
path: root/src/db.py
blob: f48960d60f5cb5d51075c90b143c96208506dea4 (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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
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,
                voting_message_ts TEXT,
                FOREIGN KEY (deployment_id) REFERENCES Deployments(id) ON DELETE CASCADE,
                FOREIGN KEY (user_id) REFERENCES Users(slack_id)
             )
         """)
        await db.execute("""
             CREATE TABLE IF NOT EXISTS Votes (
                deployment_id INTEGER NOT NULL,
                voter_id TEXT NOT NULL,
                vote TEXT NOT NULL,
                voted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (deployment_id) REFERENCES Deployments(id) ON DELETE CASCADE,
                FOREIGN KEY (voter_id) REFERENCES Users(slack_id)
                PRIMARY KEY (deployment_id, voter_id)
             )
         """)
        await db.execute("""
             CREATE TABLE IF NOT EXISTS ChannelApprovals (
                channel_id TEXT PRIMARY KEY,
                approved_by TEXT NOT NULL,
                approved_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (approved_by) 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]

async def add_vote(deployment_id: int, voter_id: str, vote: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("INSERT OR REPLACE INTO Votes (deployment_id, voter_id, vote) VALUES (?, ?, ?)", (deployment_id, voter_id, vote))
        await db.commit()

async def get_votes(deployment_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute("SELECT voter_id, vote FROM Votes WHERE deployment_id = ?", (deployment_id,))
        rows = await cursor.fetchall()
        
        votes = {"approve": [], "reject": []}
        for row in rows:
            votes[row["vote"]].append(row["voter_id"])
        return votes

async def get_user_vote(deployment_id: int, voter_id: str):
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute("SELECT vote from Votes WHERE deployment_id = ? AND voter_id = ?", (deployment_id, voter_id))
        row = await cursor.fetchone()
        return row[0] if row else None

async def get_vote_counts(deployment_id: int):
    votes = await get_votes(deployment_id)
    return len(votes["approve"]), len(votes["reject"])

async def update_voting_message_ts(deployment_id: int, message_ts: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE Deployments SET voting_message_ts = ? WHERE id = ?",
            (message_ts, deployment_id)
        )
        await db.commit()

async def approve_channel(channel_id: str, approver_id: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR REPLACE INTO ChannelApprovals (channel_id, approved_by) VALUES (?, ?)",
            (channel_id, approver_id)
        )
        await db.commit()

async def is_channel_approved(channel_id: str) -> bool:
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(
            "SELECT channel_id FROM ChannelApprovals WHERE channel_id = ?",
            (channel_id,)
        )
        row = await cursor.fetchone()
        return row is not None

async def get_approved_channels() -> list[str]:
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute("SELECT channel_id FROM ChannelApprovals")
        rows = await cursor.fetchall()
        return [row[0] for row in rows]

async def revoke_channel_approval(channel_id: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "DELETE FROM ChannelApprovals WHERE channel_id = ?",
            (channel_id,)
        )
        await db.commit()