diff --git a/config.json b/config.json index 01bbaed..90a0e29 100644 --- a/config.json +++ b/config.json @@ -14,5 +14,8 @@ }, "Logging": { "Console": true + }, + "Database": { + "ConnectionString": "mysql://user:password@host:port/database" } } diff --git a/config.yaml b/config.yaml index f325370..041bd69 100644 --- a/config.yaml +++ b/config.yaml @@ -10,4 +10,6 @@ SASL: SASLNick: EliteBot SASLPassword: password Logging: - Console: true \ No newline at end of file + Console: true +Database: + ConnectionString: mysql://user:password@host:port/database \ No newline at end of file diff --git a/plugins/cookie.py b/plugins/cookie.py new file mode 100644 index 0000000..f8170af --- /dev/null +++ b/plugins/cookie.py @@ -0,0 +1,58 @@ +from sqlalchemy import Table, Column, Integer, String, MetaData, insert, select + +from src.channel_manager import ChannelManager +from src.db import Database +from src.plugin_base import PluginBase + +meta = MetaData() +cookie_table = Table( + 'Cookie', + meta, + Column('id', Integer, primary_key=True, autoincrement=True), + Column('name', String, unique=True, nullable=False), + Column('cookies', Integer, default=0), + Column('last', String, default='1999/01/01 00:00:00'), +) +c_db = Database(cookie_table, meta) + + +class Plugin(PluginBase): + def handle_message(self, source_nick, channel, message): + parts = message.split() + c_db.create_table('Cookie') + self.channel_manager = ChannelManager() + + if parts[0].lower() == '!cookie': + if len(parts) == 1: # !cookie + self.insert_user(source_nick) + self.bot.ircsend(f'PRIVMSG {channel} :Nooooo~~') + + c_db.set(source_nick, {'cookies': 1, 'last': '1999/01/01 00:00:01'}) + elif len(parts) == 2: # !cookie USER + cookies = c_db.get(parts[1], 2) + + if cookies == -1: + self.bot.ircsend(f'PRIVMSG {channel} :I\'ve looked everywhere for {parts[1]}, but I couldn\'t ' + f'find them.') + else: + c = 'cookie' + if cookies == 0: + c = 'no cookies.' + elif cookies == 1: + c = f'{cookies} cookie.' + else: + c = f'{cookies} cookies.' + + self.bot.ircsend(f'PRIVMSG {channel} :{parts[1]} currently has {c}') + + def insert_user(self, user: str): + with c_db.engine.connect() as conn: + stmt = select(cookie_table).where(cookie_table.c.name == user) + cnt = len(conn.execute(stmt).fetchall()) + + if cnt == 0: + conn.execute(( + insert(cookie_table). + values({'name': user}) + )) + conn.commit() diff --git a/requirements.txt b/requirements.txt index d0fd917..a901ff7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,7 @@ colorama==0.4.6 +greenlet==3.0.3 +mysql-connector==2.2.9 PyYAML==6.0.1 +SQLAlchemy==2.0.27 +SQLAlchemy-Utils==0.41.1 +typing_extensions==4.9.0 diff --git a/src/bot.py b/src/bot.py index e1ef083..d34cb19 100644 --- a/src/bot.py +++ b/src/bot.py @@ -21,6 +21,7 @@ class Bot: def __init__(self, config_file): self.config = self.load_config(config_file) self.validate_config(self.config) + self.connection_string = self.config['Database'].get('ConnectionString') self.channel_manager = ChannelManager() self.logger = Logger('logs/elitebot.log') self.connected = False @@ -36,7 +37,8 @@ class Bot: ["Connection", "Nick"], ["Connection", "Ident"], ["Connection", "Name"], - ["SASL", "UseSASL"] + # ["SASL", "UseSASL"], + ["Database", "ConnectionString"] ] for field in required_fields: diff --git a/src/db.py b/src/db.py new file mode 100644 index 0000000..2e252b0 --- /dev/null +++ b/src/db.py @@ -0,0 +1,39 @@ +from sqlalchemy import create_engine, Table, MetaData, inspect, update, select +from sqlalchemy_utils import database_exists, create_database + + +# TODO: Load ConnectionString from config +class Database: + def __init__(self, table: Table, meta: MetaData): + self.engine = create_engine(r'sqlite:///.\data\database.db') + self.table = table + self.meta = meta + + if not database_exists(self.engine.url): + create_database(self.engine.url) + + def create_table(self, table_name): + if not inspect(self.engine).has_table(table_name): + self.meta.create_all(self.engine) + + def set(self, user: str, values: dict): + with self.engine.connect() as conn: + stmt = select(self.table).where(self.table.c.name == user) + cnt = len(conn.execute(stmt).fetchall()) + + if cnt == 1: + conn.execute(( + update(self.table). + values(values) + )) + conn.commit() + + def get(self, user: str, index: int): + with self.engine.connect() as conn: + stmt = select(self.table).where(self.table.c.name == user) + cnt = len(conn.execute(stmt).fetchall()) + + if cnt == 1: + return conn.execute(select(self.table).where(self.table.c.name == user)).fetchone()[index] + else: + return -1