-
Notifications
You must be signed in to change notification settings - Fork 0
/
bot.py
157 lines (135 loc) · 5.61 KB
/
bot.py
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
# -*- coding: utf-8 -*-
""" Discord Bot Main.
This module contains the Deskom bot class that takes an optional token argument.
Usage:
$ python3 bot.py --token <your-token-here>
"""
import argparse
import sqlite3
from os import getenv
from sqlite3 import Cursor
from discord import Embed
from discord.ext import commands
from discord.ext import tasks
from dotenv import load_dotenv
from emoji import EMOJI_UNICODE as EMOJIS
from cogs.eskom_cog import EskomCog
from core.eskom_interface import EskomInterface
load_dotenv()
class Deskom:
""" Eskom Bot Class """
def __init__(self, token: str = None):
# setup bot
self.token = token if token else getenv('DISCORD_TOKEN')
if not self.token:
raise Exception('Token missing. Make sure token is set in your env or passed as an argument.')
self.cogs = [
{'name': 'Eskom', 'obj': EskomCog, 'active': True},
]
self.bot = commands.Bot(command_prefix='!eskom ', case_insensitive=False)
# database connection
self.con = sqlite3.connect('eskom.sqlite.db')
self.setup_tables()
# access to eskom interface
self.eskom_interface = EskomInterface()
# assign connection to bot
setattr(self.bot, 'con', self.con)
setattr(self.bot, 'eskom_interface', self.eskom_interface)
# init functions
self.init_cogs()
self.add_events()
def setup_tables(self):
""" Create `channels` and `eskom_stage` tables if not already created """
cursor = self.con.cursor()
cursor.execute(
'CREATE TABLE IF NOT EXISTS '
'channels('
' id integer PRIMARY KEY,'
' guild_name TEXT NOT NULL,'
' guild_id TEXT NOT NULL UNIQUE,'
' channel_name TEXT NOT NULL,'
' channel_id TEXT NOT NULL UNIQUE,'
' announce INTEGER DEFAULT 1 NOT NULL'
')'
)
cursor.execute(
'CREATE TABLE IF NOT EXISTS '
'eskom_stage('
' id integer PRIMARY KEY,'
' stage INTEGER NOT NULL'
')'
)
def create_announcement_embed(self, new_stage: int) -> Embed:
"""
Create announcement embed based on current stage difference to last stage recorded.
:param new_stage: New stage found.
:return: Embed object with correct announcement message and color.
"""
cursor = self.con.cursor()
cursor.execute('SELECT stage, created_at from eskom_stage ORDER BY created_at DESC LIMIT 1')
last_stage = cursor.fetchone()[0] if cursor.fetchone() else 0
cursor.close()
stage_change_colors = {'good': 0x2bcc2b, 'warning': 0xccaf2b, 'bad': 0xcc622b}
if new_stage == 0:
color = stage_change_colors['good']
message = f'We have moved down to stage {new_stage}.'
elif new_stage < last_stage:
color = stage_change_colors['warning']
message = f'We moved down to stage {new_stage}. We aren\'t in the clear just yet guys.'
else:
color = stage_change_colors['bad']
message = f'We moved up to stage {new_stage}. Get your candles ready and phones charged.'
embed_message = Embed(
title=f"{EMOJIS['en'][':flashlight:']} Eskom Stage Change Announcement!",
description=message,
color=color
)
return embed_message
async def announce_stage_change(self, cursor: Cursor, stage: str):
"""
Query all channels to announce. Loop through channels and announce new stage one at a time.
:param cursor: database cursor
:param stage: current load-shedding stage to announce
"""
cursor.execute('SELECT channel_id FROM channels WHERE announce = "1"')
rows = cursor.fetchall()
announce_embed = self.create_announcement_embed(new_stage=int(stage))
for row in rows:
channel = self.bot.get_channel(int(row[0]))
await channel.send(embed=announce_embed)
self.con.commit()
@tasks.loop(seconds=60)
async def lookup_eskom_stage(self):
""" Long running looping task that queries Eskom page for latest stage and announces if changed """
await self.bot.wait_until_ready()
cursor = self.con.cursor()
stage = await self.eskom_interface.async_get_stage()
cursor.execute('SELECT stage FROM eskom_stage ORDER BY ROWID DESC LIMIT 1')
old_stage = cursor.fetchone()
old_stage = old_stage[0] if old_stage else old_stage
if stage != old_stage:
cursor.execute('INSERT INTO eskom_stage(stage) VALUES (?)', (stage,))
self.con.commit()
await self.announce_stage_change(cursor=cursor, stage=str(stage))
cursor.close()
def init_cogs(self):
""" Add cogs to bot """
for cog in self.cogs:
self.bot.add_cog(cog['obj'](self.bot))
async def on_ready(self):
""" Print to console once bot is connected """
print(f'{self.bot.user.name} is connected.')
def add_events(self):
""" Add events to bot """
self.bot.event(self.on_ready)
def start_bot(self):
""" Run long running loop task and then start up the bot """
# pylint: disable=E1101
self.lookup_eskom_stage.start()
self.bot.run(self.token)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Eskom discord bot.')
parser.add_argument('-t', '--token', type=str, help='Use discord token for bot.')
args = parser.parse_args()
client = Deskom(token=args.token)
client.start_bot()