From 12b667b58e271e98e899420071337e99463d1bb9 Mon Sep 17 00:00:00 2001
From: liyunze <50455574+liyunze-coding@users.noreply.github.com>
Date: Fri, 3 May 2024 14:32:08 +0800
Subject: [PATCH] fixed critical bug, added kill switch command for mods and
 streamer

---
 functions.py | 38 +++++++++++++++++++++++++++-----------
 main.py      |  7 +++++--
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/functions.py b/functions.py
index 78db124..1c9969e 100644
--- a/functions.py
+++ b/functions.py
@@ -9,12 +9,13 @@ import win32con  # Windows only
 
 
 class TwitchPlaysX:
-    def __init__(self, key_delay, keymap, channel_name, x_limit) -> None:
+    def __init__(self, key_delay, keymap, channel_name, x_limit, kill_commands) -> None:
         self.os = self.get_os()
         self.key_delay = key_delay
         self.keymap = keymap
         self.channel_name = channel_name.lower()
         self.x_limit = x_limit
+        self.kill_commands = kill_commands
 
     # not tested
     def sendKeyLinux(self, button):
@@ -51,27 +52,36 @@ class TwitchPlaysX:
             match = re.match(r"([a-z]+)(\d*)", char)
             if match:
                 control, count = match.groups()
+                if control not in self.keymap:
+                    return []
                 count = int(count or "1")
                 if count > self.x_limit:
                     count = self.x_limit
                 compiled_chars.extend([control] * count)
+
         return compiled_chars
 
-    def process_message(self, author, message):
+    def process_message(self, author, message, is_mod):
+        if is_mod and (message.lower() in self.kill_commands):
+            print("Killing TwitchPlaysX")
+            exit()
+
         characters = message.lower().split(" ")
 
         os_type = self.get_os()
 
-        if self.all_valid_chars(characters):
-            compiled_characters = self.compile_chars(characters)
-            for char in compiled_characters:
-                print(f"{author}: {char}")
-                self.sendKey(char, os_type)
+        compiled_characters = self.compile_chars(characters)
+        for char in compiled_characters:
+            print(f"{author}: {char}")
+            self.sendKey(char, os_type)
 
-    async def listen(self):
+    async def play(self):
         reader, writer = await asyncio.open_connection("irc.chat.twitch.tv", 6667)
 
         writer.write(f"NICK justinfan1234\r\n".encode("utf-8"))
+        writer.write(
+            f"CAP REQ :twitch.tv/tags\r\n".encode("utf-8")
+        )  # Enable IRCv3 tags
         writer.write(f"JOIN #{self.channel_name}\r\n".encode("utf-8"))
 
         while True:
@@ -83,10 +93,16 @@ class TwitchPlaysX:
                 await writer.drain()
 
             elif message:
-                author = message.split("!", 1)[0].split(":", 1)[-1]
+                tags, author, msg = message.split(" ", 2)
+                author = author.split("!", 1)[0].split(":", 1)[-1]
                 msg = (
-                    message.split(f"#{self.channel_name.lower()}", 1)[-1]
+                    msg.split(f"#{self.channel_name.lower()}", 1)[-1]
                     .split(":", 1)[-1]
                     .strip()
                 )
-                self.process_message(author, msg)
+
+                is_mod = ("user-type=mod" in tags) or (
+                    "badges=broadcaster" in tags
+                )  # Check if the user is a moderator
+
+                self.process_message(author, msg, is_mod)
diff --git a/main.py b/main.py
index 315ee4a..ca27e6a 100644
--- a/main.py
+++ b/main.py
@@ -7,6 +7,7 @@ import asyncio
 CHANNEL_NAME = "RythonDev"
 KEYDELAY = 0.3
 X_LIMIT = 10  # "up20" will be converted to "up10" if X_LIMIT is 10
+KILL_SWITCH_COMMANDS = ["!killpokemon"]  # must be all lowercase
 
 # Keymap for Windows
 windowsKeyMap = {
@@ -46,10 +47,12 @@ async def main():
 
         keymap = windowsKeyMap if os_type == "Windows" else linuxKeyMap
 
-        twitchPlaysX = TwitchPlaysX(KEYDELAY, keymap, CHANNEL_NAME, X_LIMIT)
+        twitchPlaysX = TwitchPlaysX(
+            KEYDELAY, keymap, CHANNEL_NAME, X_LIMIT, KILL_SWITCH_COMMANDS
+        )
 
         print(f"Starting TwitchPlaysX for channel: {CHANNEL_NAME}")
-        await twitchPlaysX.listen()
+        await twitchPlaysX.play()
 
 
 if __name__ == "__main__":