Prepend PING messages with "proxy:"

1 files changed, 18 insertions, 6 deletions

diff --git a/idc_irc_proxy.py b/idc_irc_proxy.py
index c02573b..dab6380 100755
--- a/idc_irc_proxy.py
+++ b/idc_irc_proxy.py
@@ -39,6 +39,8 @@ class Proxy:
 
     # Send messages
     def send(self, cmd, hostmask, tags, args):
+        if hostmask == ('', '', ''):
+            hostmask = None
         raw = self._irc_msg_unparser(
             cmd, hostmask or (cmd, cmd, cmd), tags, args, colon=False,
             encoding=self.encoding
@@ -79,8 +81,11 @@ class Proxy:
             threading.Thread(target=self._main).start()
         elif cmd == 'ERROR':
             self.send('PING', None, {}, [':ERROR'])
-        elif cmd == 'PONG' and args and args[-1] == 'miniirc-ping':
-            return
+        elif cmd == 'PONG':
+            # Translate proxied pings back
+            if not args or not args[-1].startswith('proxy:'):
+                return
+            args = [args[-1][6:]]
 
         # Send the command to the client
         try:
@@ -136,6 +141,8 @@ class Proxy:
                 self.sock.shutdown(socket.SHUT_RDWR)
                 self.sock.close()
                 return
+            elif cmd == 'PING':
+                self.send('PONG', None, {}, args)
             else:
                 self._sendq.append((tags, cmd, args))
 
@@ -149,10 +156,15 @@ class Proxy:
 
     # Send a command
     def _sendcmd(self, tags, cmd, args):
-        if cmd not in self.block_outgoing:
-            raw = ircv2_message_unparser(cmd, (cmd, cmd, cmd), {}, args,
-                colon=False, encoding=None)
-            self.irc.quote(raw, tags=tags)
+        if cmd in self.block_outgoing:
+            return
+
+        if cmd == 'PING':
+            args = [f'proxy:{args[0] if args else ""}']
+
+        raw = ircv2_message_unparser(cmd, (cmd, cmd, cmd), {}, args,
+                                     colon=False, encoding=None)
+        self.irc.quote(raw, tags=tags)
 
     # The more permanent main loop
     def _main(self, single_thread=False):