diff --git a/main.py b/main.py index e0d2def..77000d6 100755 --- a/main.py +++ b/main.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import os +import sys import socket import signal import threading @@ -7,15 +8,23 @@ import subprocess import re import time import queue +import config +import ssl as ssllib class Server(): + # Order of functions in Server(): + # - Public API (non underscored functions and __init__) + # - Thread targets + # - Internal helper functions (message parsing, message handling, etc) + # - Bot commands (pid, die, etc) + # Setup the IRC related things such as user details and prefixes # This function also sets up the threading events and queues - def __init__(self, realname, nickname, channel, command_prefix = "$!", bot_prefix = "$$", opper_nicknames = []): + def __init__(self, realname, nickname, channels, command_prefix = "$!", bot_prefix = "$$", opper_nicknames = [], message_queue_max_size = 512): print(f"[SERVER/MAINTHREAD] Using Server() on Python3 PID {os.getpid()}") self.realname = realname self.nickname = nickname - self.channel = channel + self.channels = channels self.opper_nicknames = opper_nicknames # Prefix for commands to execute @@ -27,21 +36,33 @@ class Server(): # Used to signal threads to shutdown self._going_down = threading.Event() - # Queue used for messages to send - self._msg_q = queue.Queue() + # Used to signal command thread to kill the children process + self._kill_cmd = threading.Event() + + # Queue used for IRC server send() calls + self._send_q = queue.Queue(maxsize = message_queue_max_size) # Rate limiting variables # Used so it can be quired by IRC self._msg_time = 0 self._msg_count = 0 - print(f"[SERVER/MAINTHREAD] Using nickname {nickname} on {channel}!") + print(f"[SERVER/MAINTHREAD] Using nickname {nickname}!") # This function handles the socket bring up # Sets socket paramaters, and starts the send/recv threads - def connect(self, ip, port, sock_timeout = 60, sock_sendbuf = 512, sock_recvbuf = 512): + def connect(self, ip, port, ssl = False, sock_timeout = 60, sock_sendbuf = 512, sock_recvbuf = 512): print(f"[SERVER/MAINTHREAD] Connecting to {ip}:{port}") - self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + + if ssl: + print(f"[SERVER/MAINTHREAD] Using SSL for connection!") + context = ssllib.create_default_context() + self.sock = context.wrap_socket(sock, server_hostname = ip) + else: + print(f"[SERVER/MAINTHREAD] Using plain text for connection!") + self.sock = sock print(f"[SERVER/MAINTHREAD] Using socket timeout of {sock_timeout} seconds!") self.sock.settimeout(sock_sendbuf) @@ -60,40 +81,72 @@ class Server(): # Gracefully shuts the server down def die(self, msg = "Python3 Server() outta here!"): - print("[SERVER/MAINTHREAD] Sending PART and QUIT") - # Part all of our channels with a cool message, then quit - self._msg_q.put(f"PART {self.channel} : {msg}\r\n".encode()) - self._msg_q.put(f"QUIT : {msg}\r\n".encode()) + # NOOP if we're in the process of going down + if self._going_down.is_set(): + return - print("[SERVER/MAINTHREAD] Signaling threads to quit!") + print("[SERVER] Sending PART and QUIT") + # Part all of our channels with a cool message, then quit + for channel in self.channels: + self.sock.send(f"PART {channel} : {msg}\r\n".encode()) + + self.sock.send(f"QUIT : {msg}\r\n".encode()) + + print("[SERVER] Signaling threads to quit!") # Signal threads to quit self._going_down.set() + print("[SERVER] Signaling main thread to quit!") + os.kill(os.getpid(), signal.SIGINT) + + # Sends a message to a target (user/channel) + # The bot does NOT have to be a channel to send a message, but most channels + # disallow messages from non-JOIN'ed users. + def privmsg(self, target, message, bypass_q = False): + # Check for send thread + try: + self._send_thread + except NameError: + raise RuntimeWarning("connect() must be called before this command is used!") + + # format message + message = f"PRIVMSG {target} :{message}\r\n".encode() + + if bypass_q: + self.sock.send(message) + else: + self._send_q.put(message) + # Function used to send messages from the queue # Rate limiting is also implemented here def _send_loop(self, ip): while not self._going_down.is_set(): - self._msg_count+=1 self._msg_time = (0.15 * self._msg_count)**2 print(f"[SENDTHREAD] Sleeping for {self._msg_time} seconds on message count {self._msg_count}") time.sleep(self._msg_time) if self._msg_count == 10: + print("[SENDTHREAD] Sending ping on 10th message!") + self.sock.send(f"PING :{ip}\r\n".encode()) self._msg_count = 0 # Grab message and send it try: - msg = self._msg_q.get(timeout = 60) + msg = self._send_q.get(timeout = 60) + + # Only increment the message count if we actually get a message, + # the server will PING us when it needs to. + self._msg_count+=1 except queue.Empty: # So the thread signaler works - print("[SENDTHREAD] Timeout occurred waiting on queue... sending ping!") - self.sock.send(f"PING :{ip}\r\n".encode()) + continue print(f"[SENDTHREAD] Sending message: {msg}") try: self.sock.send(msg) except BrokenPipeError: # Can't use die because we have no connection... just quit I suppose - self._going_down.set() + print("[SENDTHREAD] Quitting due to BrokenPipeError!") + break print("[SENDTHREAD] Quitting due to thread condition!") @@ -120,10 +173,7 @@ class Server(): print("[RECVTHREAD] Message: {data}") continue except socket.timeout: - # Instead of using another thread, I just use a socket timeout to send pings - # Kills two birds with one stone as we can't wait forever due to threading Events - print("[RECVTHREAD] Timeout occurred waiting on server messages... sending ping!") - self.sock.send(f"PING :{ip}\r\n".encode()) + # So threading Events work continue # Handle IRC messages @@ -139,13 +189,67 @@ class Server(): if self._going_down.is_set(): print("[RECVTHREAD] Quitting due to thread condition!") + # This is where we actually run the RCE commands and pipe the output + # back to IRC. + # + # No temp files here + def _handle_command(self, cmd, target_channel): + print(f"[CMDTHREAD] Running CMD {cmd}!") + proc = subprocess.Popen(cmd, shell = True, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, stdin = subprocess.DEVNULL, + start_new_session=True, text=False) + + # Calculate maximum message size + # 512 bytes is max IRC message with