From a4f97db581e45fb2603ceda237e3aefe07df2074 Mon Sep 17 00:00:00 2001
From: Ian Mancini <ianmethyst@gmail.com>
Date: Wed, 13 May 2020 20:17:10 -0300
Subject: [PATCH] Improve twitch notifications script

---
 launch-twitch-notify |   6 ++-
 setup/setup_pacman   |   6 +--
 twitch-notify        | 107 ++++++++++++++++++++++++++++++-------------
 3 files changed, 84 insertions(+), 35 deletions(-)

diff --git a/launch-twitch-notify b/launch-twitch-notify
index ab0c551..79f6285 100755
--- a/launch-twitch-notify
+++ b/launch-twitch-notify
@@ -1,3 +1,7 @@
 #!/bin/bash
 
-TWITCH_OAUTH=$(gpg -q --decrypt data/twitch-oauth.gpg) python3 twitch-notify
+SCRIPT_PATH=$(dirname "$0")
+TOKEN_PATH="$SCRIPT_PATH/data/twitch-oauth.gpg"
+
+TWITCH_OAUTH=$(gpg -q --decrypt $TOKEN_PATH ) \
+  exec python3 "$SCRIPT_PATH/twitch-notify"
diff --git a/setup/setup_pacman b/setup/setup_pacman
index 3e23565..1eab864 100755
--- a/setup/setup_pacman
+++ b/setup/setup_pacman
@@ -16,7 +16,7 @@ if [[ $enablewget =~ [yY] ]]; then
   pacman -Q wget &> /dev/null
   if [ $? = 1 ]; then
     echo ":: Installing wget"
-    pacman -Syu wget 
+    pacman -Syu wget
   else
     echo ":: Success: wget is installed"
   fi
@@ -33,8 +33,8 @@ if [[ $enablemultilib =~ [yY] ]]; then
   echo "Enabling multilib"
 
   if ! grep -q "^\[multilib\]" /etc/pacman.conf; then
-    cat >> /etc/pacman.conf <<EOF 
-  
+    cat >> /etc/pacman.conf <<EOF
+
 [multilib]
 Include = /etc/pacman.d/mirrorlist
 EOF
diff --git a/twitch-notify b/twitch-notify
index 0eccc2d..1e603cb 100644
--- a/twitch-notify
+++ b/twitch-notify
@@ -1,9 +1,12 @@
 #!/usr/bin/python3
 
-from subprocess import call
+from subprocess import run
+from textwrap import dedent
 import os
+import sys
 import socket
 import re
+import signal
 
 config = {
     "token": None,
@@ -13,54 +16,96 @@ config = {
     "short_notification": True
 }
 
-
-try:
-    config['token'] = os.environ['TWITCH_OAUTH']
-except KeyError:
-    print("Get your OAUTH token here: https://twitchapps.com/tmi/",
-          "\nThen pass it as an environment variable:",
-          f"\nTWITCH_OAUTH=<token> {__file__}")
-    exit(1)
-
-
 HOST = 'irc.chat.twitch.tv'
 PORT = 6667
-MESSAGE_REGEX = ':(.*)\!.*@.* PRIVMSG #(.*) :(.*)'
-
+MESSAGE_REGEX = r':(.*)!.*@.* PRIVMSG #(.*) :(.*)'
 
 sock = socket.socket()
 connected = False
 
 
+def get_token_from_env():
+    """Get OAUTH token from environment variables or exit"""
+
+    try:
+        config['token'] = os.environ['TWITCH_OAUTH']
+    except KeyError:
+        helpText = """
+                Get your OAUTH token here: https://twitchapps.com/tmi/
+                Then pass it as an environment variable:
+                TWITCH_OAUTH=<token> {__file__}"
+                """
+        print(dedent(helpText))
+        sys.exit(1)
+
+
 def connect():
+    """Connect to IRC and authenticate"""
+
     sock.connect((HOST, PORT))
     sock.send(f"PASS {config['token']}\n".encode('utf-8'))
     sock.send(f"NICK {config['nickname']}\n".encode('utf-8'))
     sock.send(f"JOIN {config['channel']}\n".encode('utf-8'))
 
 
-connect()
-
-while True:
-    resp = sock.recv(2048).decode('utf-8')
+def pong(resp):
+    """Answer to server PING to keep socket connection alive"""
 
     if resp.startswith('PING'):
         sock.send("PONG\n".encode('utf-8'))
+        return True
+    return False
+
+
+def wait_for_chat(resp):
+    """Wait until all welcome and user list messages are received"""
+
+    if "End of /NAMES list" in resp:
+        global connected
+        connected = True
+        run(["notify-send",
+            f"connected to {config['channel']} as {config['nickname']}"])
+
+
+def handle_message(resp):
+    """Parse the message and send it as a notification"""
+
+    print(resp)
+    user, channel, message = re.search(MESSAGE_REGEX, resp).groups()
+
+    if user and channel and message:
+        if config['short_notification']:
+            notification_text = f"@{user}: {message}"
+        else:
+            notification_text = f"@{user} in #{channel}\n{message}"
+
+        if not (config['supress_own'] and config['nickname'] == user):
+            run(["notify-send", notification_text])
+
+
+def handle_exit(sig, frame):
+    """Close socket connection before exit"""
+
+    print(f"Got {sig.name}, closing socket")
+    sock.shutdown(socket.SHUT_RDWR)
+    sock.close()
+    sys.exit(0)
+
 
-    elif len(resp) > 0:
-        if connected:
-            user, channel, message = re.search(MESSAGE_REGEX, resp).groups()
+if __name__ == "__main__":
+    # Register signal handlers
+    signal.signal(signal.SIGINT, handle_exit)
+    signal.signal(signal.SIGTERM, handle_exit)
 
-            if config['short_notification']:
-                notification_text = f"@{user}: {message}"
-            else:
-                notification_text = f"@{user} in #{channel}\n{message}"
+    get_token_from_env()
+    connect()
 
-            if not (config['supress_own'] and config['nickname'] == user):
-                call(["notify-send", notification_text])
+    while True:
+        resp = sock.recv(2048).decode('utf-8')
 
-        if "End of /NAMES list" in resp:
-            connected = True
-            call([
-                "notify-send",
-                f"connected to {config['channel']} as {config['nickname']}"])
+        if len(resp) > 0:
+            if not (pong(resp)):
+                if connected:
+                    handle_message(resp)
+                else:
+                    wait_for_chat(resp)