added todos as well as regex based command triggers.HEADmaster

9 files changed, 207 insertions, 69 deletions

diff --git a/commands/__init__.py b/commands/__init__.py
index dc26d49..a7b0091 100644
--- a/commands/__init__.py
+++ b/commands/__init__.py
@@ -1,2 +1 @@
-from commands.handling import handle_commands, load_commands
-
+from commands.handling import handle_commands, load_modules, handle_match_scripts
diff --git a/commands/handling.py b/commands/handling.py
index 46cdb66..9c9e864 100644
--- a/commands/handling.py
+++ b/commands/handling.py
@@ -19,9 +19,11 @@ def load_module(module):
     for func in functions:
         if lib.is_command(func):
             lib.register_command(func)
+        if lib.is_match_script(func):
+            lib.register_match_script(func)
 
 
-def load_commands(folder='modules'):
+def load_modules(folder='modules'):
     for dirname, dirnames, filenames in os.walk(folder):
         for filename in filenames:
             filename: str
@@ -32,10 +34,14 @@ def load_commands(folder='modules'):
                 load_module(module)
 
 
-def handle_commands(client: pyrogram.Client, update, users, chats):
-    if not (isinstance(update, tgtypes.UpdateNewMessage)
+def is_message_update(update):
+    return (isinstance(update, tgtypes.UpdateNewMessage)
             or isinstance(update, tgtypes.UpdateNewChannelMessage)
-            or isinstance(update, tgtypes.UpdateNewEncryptedMessage)):
+            or isinstance(update, tgtypes.UpdateNewEncryptedMessage))
+
+
+def handle_commands(client: pyrogram.Client, update, users, chats):
+    if not is_message_update(update):
         return
     update: tgtypes.UpdateNewMessage
     message: tgtypes.Message = update.message
@@ -51,7 +57,7 @@ def handle_commands(client: pyrogram.Client, update, users, chats):
         return
     command = parts[0][1:]
     args = parts[1:]
-    cmd_func = lib.commands[command.lower()]
+    cmd_func = lib.get_command_by_name(command.lower())
     ctx = lib.CommandContext(client=client, channel=message.to_id, args=args, message=message)
     try:
         cmd_func(ctx)
@@ -60,3 +66,22 @@ def handle_commands(client: pyrogram.Client, update, users, chats):
     except Exception as e:
         ctx.respond("unknown exception during execution. Error will be DM'd" + str(e))
         print(traceback.format_exc(), file=sys.stderr)
+
+
+def handle_match_scripts(client: pyrogram.Client, update, users, chats):
+    if not is_message_update(update):
+        return
+    update: tgtypes.UpdateNewMessage
+    message: tgtypes.Message = update.message
+    author_id = message.from_id
+    if author_id != client.user_id:
+        # do not react to other people
+        return
+    text: str = message.message
+    for regex, func in lib.get_match_scripts().items():
+        match = regex.match(text)
+        if match is None:
+            continue
+        ctx = lib.MatchContext(client=client, channel=message.to_id, message=message,
+                               match=match, groups=match.groups(), named_groups=match.groupdict())
+        func(ctx)
diff --git a/lib/__init__.py b/lib/__init__.py
index edb1b42..0cec7cf 100644
--- a/lib/__init__.py
+++ b/lib/__init__.py
@@ -1,61 +1,5 @@
-from typing import List
+from .commands import CommandContext, get_all_commands, get_command_name, get_command_description, description, name, \
+    is_command, register_command, get_command_by_name
 
-import pyrogram
-from pyrogram.api import types as tgtypes
-
-
-def property_decorator(key):
-    def decorator(value):
-        def wrapper(func):
-            setattr(func, key, value)
-            return func
-
-        return wrapper
-
-    return decorator
-
-
-name = property_decorator('name')
-description = property_decorator('description')
-
-
-def is_command(func):
-    return hasattr(func, 'name') and func.name is not None
-
-
-def get_command_name(func):
-    return func.name
-
-
-def get_command_description(func):
-    return func.description
-
-
-commands = {}
-
-
-def get_all_commands():
-    return commands.values()
-
-
-class CommandContext(object):
-    def __init__(self, client: pyrogram.Client, channel, args: List[str], message: tgtypes.Message):
-        import re
-        self.args = args
-        self.client = client
-        self.channel = channel
-        self.message = message
-        self.rest_content = re.sub('^.*? ', '', message.message)
-        self.author = message.from_id
-
-    def respond(self, text):
-        self.client.send_message(self.channel, text=text)
-
-    def edit(self, text):
-        self.client.edit_message_text(chat_id=self.channel, message_id=self.message.id, text=text)
-
-
-def register_command(func):
-    if not is_command(func):
-        return
-    commands[get_command_name(func)] = func
+from .match_script import get_match_script_matcher, is_match_script, match_text, register_match_script, \
+    get_match_scripts, MatchContext
diff --git a/lib/commands.py b/lib/commands.py
new file mode 100644
index 0000000..638ab9b
--- /dev/null
+++ b/lib/commands.py
@@ -0,0 +1,48 @@
+import typing
+
+import pyrogram
+from pyrogram.api import types as tgtypes
+
+from lib.common import CommonContext
+from util import property_decorator
+
+commands = {}
+
+
+def get_all_commands():
+    return commands.values()
+
+
+class CommandContext(CommonContext):
+    def __init__(self, client: pyrogram.Client, channel, args: typing.List[str], message: tgtypes.Message):
+        super().__init__(client, channel, message)
+        import re
+        self.args = args
+        self.rest_content = re.sub('^.*? ', '', message.message)
+
+
+
+def register_command(func):
+    if not is_command(func):
+        return
+    commands[get_command_name(func)] = func
+
+
+def is_command(func):
+    return hasattr(func, 'name') and func.name is not None
+
+
+def get_command_name(func):
+    return func.name
+
+
+def get_command_description(func):
+    return func.description
+
+
+def get_command_by_name(command_name: str):
+    return commands[command_name.lower()]
+
+
+name = property_decorator('name')
+description = property_decorator('description')
diff --git a/lib/common.py b/lib/common.py
new file mode 100644
index 0000000..ade2891
--- /dev/null
+++ b/lib/common.py
@@ -0,0 +1,16 @@
+import pyrogram
+from pyrogram.api import types as tgtypes
+
+
+class CommonContext:
+    def __init__(self, client: pyrogram.Client, channel, message: tgtypes.Message):
+        self.client = client
+        self.channel = channel
+        self.message = message
+        self.author = message.from_id
+
+    def respond(self, text):
+        self.client.send_message(self.channel, text=text)
+
+    def edit(self, text):
+        self.client.edit_message_text(chat_id=self.channel, message_id=self.message.id, text=text)
diff --git a/lib/match_script.py b/lib/match_script.py
new file mode 100644
index 0000000..124c23b
--- /dev/null
+++ b/lib/match_script.py
@@ -0,0 +1,42 @@
+import pyrogram
+from pyrogram.api import types as tgtypes
+
+from lib.common import CommonContext
+
+
+def match_text(regex):
+    import re
+    regex = re.compile(regex)
+
+    def wrapper(func):
+        func.regex = regex
+        return func
+
+    return wrapper
+
+
+match_scripts = {}
+
+
+def is_match_script(func):
+    return hasattr(func, 'regex') and func.regex is not None
+
+
+def get_match_script_matcher(func):
+    return func.regex
+
+
+def register_match_script(func):
+    match_scripts[func.regex] = func
+
+
+class MatchContext(CommonContext):
+    def __init__(self, client: pyrogram.Client, channel, message: tgtypes.Message, match, groups, named_groups):
+        super().__init__(client, channel, message)
+        self.match = match
+        self.groups = groups
+        self.named_groups = named_groups
+
+
+def get_match_scripts():
+    return match_scripts
diff --git a/main.py b/main.py
index 693f9d4..ebfec76 100644
--- a/main.py
+++ b/main.py
@@ -1,14 +1,15 @@
 from pyrogram import Client
 
-from commands import handle_commands, load_commands
+from commands import handle_commands, load_modules, handle_match_scripts
 
 
 def update_handler(client, update, users, chats):
     handle_commands(client, update, users, chats)
+    handle_match_scripts(client, update, users, chats)
 
 
 def main():
-    load_commands()
+    load_modules()
     client = Client(session_name="userbot")
     client.set_update_handler(update_handler)
     client.start()
diff --git a/modules/todo/todo.py b/modules/todo/todo.py
new file mode 100644
index 0000000..89a2b78
--- /dev/null
+++ b/modules/todo/todo.py
@@ -0,0 +1,52 @@
+import sqlite3
+
+import lib
+
+
+def init_todo_db(conn: sqlite3.Connection):
+    c = conn.cursor()
+    c.execute(
+        """
+        CREATE TABLE IF NOT EXISTS todos (
+          id INTEGER PRIMARY KEY,
+          text VARCHAR(1024)
+        );
+        """)
+    conn.commit()
+
+
+def add_todo(todo):
+    conn = sqlite3.connect('todos.db')
+    init_todo_db(conn)
+    c = conn.cursor()
+    c.execute(
+        """
+        INSERT INTO todos (text) VALUES (?);
+        """, [todo])
+    conn.commit()
+
+
+def read_todos() -> sqlite3.Cursor:
+    conn = sqlite3.connect('todos.db')
+    init_todo_db(conn)
+    c = conn.cursor()
+    return c.execute(
+        """
+        SELECT id, text FROM todos LIMIT 50;
+        """)
+
+
+@lib.match_text('(?i)^TODO: (?P<todos>.*)$')
+def handle_todo(ctx: lib.MatchContext):
+    for todo in ctx.named_groups['todos'].split(','):
+        todo = todo.strip()
+        add_todo(todo)
+        ctx.respond(todo + ' added')
+
+
+@lib.name('todos')
+def list_todos(ctx: lib.CommandContext):
+    mes = 'Todos: \n'
+    for row in read_todos():
+        mes += '%d - %s\n' % row
+    ctx.respond(mes)
diff --git a/util.py b/util.py
new file mode 100644
index 0000000..0886bf9
--- /dev/null
+++ b/util.py
@@ -0,0 +1,11 @@
+
+def property_decorator(key):
+    def decorator(value):
+        def wrapper(func):
+            setattr(func, key, value)
+            return func
+
+        return wrapper
+
+    return decorator
+