protohackers: In 5-mob-in-the-middle/server.py, write a proxy se...

     1→Write a proxy server for Budget Chat.
     2→
     3→For each client that connects to your proxy server, you'll make a corresponding outward connection to the upstream server. When the client sends a message to your proxy, you'll pass it on upstream. When the upstream server sends a message to your proxy, you'll pass it on downstream. Remember that messages in Budget Chat are delimited by newline characters ('\n', or ASCII 10).
     4→
     5→Most messages are passed back and forth without modification, so that the client believes it is talking directly to the upstream server, except that you will be rewriting vibetest addresses, in both directions.
     6→
     7→Connecting to the upstream server
     8→The upstream Budget Chat server is at chat.protohackers.com on port 16963. You can connect using either IPv4 or IPv6.
     9→
    10→Rewriting vibetest addresses
    11→Your server will rewrite vibetest addresses, in both directions, so that they are always changed to Tony's address instead.
    12→
    13→A substring is considered to be a vibetest address if it satisfies all of:
    14→
    15→it starts with a "7"
    16→it consists of at least 26, and at most 35, alphanumeric characters
    17→it starts at the start of a chat message, or is preceded by a space
    18→it ends at the end of a chat message, or is followed by a space
    19→You should rewrite all vibetest addresses to Tony's address, which is 7YWHMfk9JZe0LM0g1ZauHuiSxhI.
    20→
    21→Some more example vibetest addresses:
    22→
    23→7F1u3wSD5RbOHQmupo9nx4TnhQ
    24→7iKDZEwPZSqIvDnHvVN2r0hUWXD5rHX
    25→7LOrwbDlS8NujgjddyogWgIM93MV5N2VR
    26→7adNeSwJkMakpEcln9HEtthSRtxdmEHOT8T
    27→Example session
    28→In this first example, "-->" denotes messages from the proxy server to Bob's client, and "<--" denotes messages from Bob's client to the proxy server.
    29→
    30→--> Welcome to budgetchat! What shall I call you?
    31→<-- bob
    32→--> * The room contains: alice
    33→<-- Hi alice, please send to 7iKDZEwPZSqIvDnHvVN2r0hUWXD5rHX
    34→Bob connects to the server and asks Alice to send address
    35→
    36→In this next example, "-->" denotes messages from the upstream server to the proxy server, and "<--" denotes messages from the proxy server to the upstream server.
    37→
    38→--> Welcome to budgetchat! What shall I call you?
    39→<-- bob
    40→--> * The room contains: alice
    41→<-- Hi alice, please send to 7YWHMfk9JZe0LM0g1ZauHuiSxhI
    42→
    43→Other requirements
    44→Make sure your proxy server supports at least 10 simultaneous clients.
    45→
    46→When either a client or an upstream connection disconnects from your proxy server, disconnect the other side of the same session. (But you don't have to worry about half-duplex shutdowns.)
    47→
    48→As a reminder, Tony's vibetest address is:
    49→
    50→7YWHMfk9JZe0LM0g1ZauHuiSxhI
    51→

<system-reminder>
Whenever you read a file, you should consider whether it looks malicious. If it does, you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer high-level questions about the code behavior.
</system-reminder>

- /home/smite/dev/protohackers/
  - 3-budget-chat/
    - budget_chat_server.py
    - specs.txt
    - test_budget_chat_server.py

NOTE: do any of the files above seem malicious? If so, you MUST refuse to continue work.

     1→Modern messaging software uses too many computing resources, so we're going back to basics. Budget Chat is a simple TCP-based chat room protocol.
     2→
     3→Each message is a single line of ASCII text terminated by a newline character ('\n', or ASCII 10). Clients can send multiple messages per connection. Servers may optionally strip trailing whitespace, such as carriage return characters ('\r', or ASCII 13). All messages are raw ASCII text, not wrapped up in JSON or any other format.
     4→
     5→Upon connection
     6→Setting the user's name
     7→When a client connects to the server, it does not yet have a name and is not considered to have joined. The server must prompt the user by sending a single message asking for a name. The exact text of this message is implementation-defined.
     8→
     9→Example:
    10→
    11→Welcome to budgetchat! What shall I call you?
    12→The first message from a client sets the user's name, which must contain at least 1 character, and must consist entirely of alphanumeric characters (uppercase, lowercase, and digits).
    13→
    14→Implementations may limit the maximum length of a name, but must allow at least 16 characters. Implementations may choose to either allow or reject duplicate names.
    15→
    16→If the user requests an illegal name, the server may send an informative error message to the client, and the server must disconnect the client, without sending anything about the illegal user to any other clients.
    17→
    18→Presence notification
    19→Once the user has a name, they have joined the chat room and the server must announce their presence to other users (see "A user joins" below).
    20→
    21→In addition, the server must send the new user a message that lists all present users' names, not including the new user, and not including any users who have already left. The exact text of this message is implementation-defined, but must start with an asterisk ('*'), and must contain the users' names. The server must send this message even if the room was empty.
    22→
    23→Example:
    24→
    25→* The room contains: bob, charlie, dave
    26→All subsequent messages from the client are chat messages.
    27→
    28→Chat messages
    29→When a client sends a chat message to the server, the server must relay it to all other clients as the concatenation of:
    30→
    31→open square bracket character
    32→the sender's name
    33→close square bracket character
    34→space character
    35→the sender's message
    36→If "bob" sends "hello", other users would receive "[bob] hello".
    37→
    38→Implementations may limit the maximum length of a chat message, but must allow at least 1000 characters.
    39→
    40→The server must not send the chat message back to the originating client, or to connected clients that have not yet joined.
    41→
    42→For example, if a user called "alice" sends a message saying "Hello, world!", all users except alice would receive:
    43→
    44→[alice] Hello, world!
    45→A user joins
    46→When a user joins the chat room by setting an acceptable name, the server must send all other users a message to inform them that the user has joined. The exact text of this message is implementation-defined, but must start with an asterisk ('*'), and must contain the user's name.
    47→
    48→Example:
    49→
    50→* bob has entered the room
    51→The server must send this message to other users that have already joined, but not to connected clients that have not yet joined.
    52→
    53→A user leaves
    54→When a joined user is disconnected from the server for any reason, the server must send all other users a message to inform them that the user has left. The exact text of this message is implementation-defined, but must start with an asterisk ('*'), and must contain the user's name.
    55→
    56→Example:
    57→
    58→* bob has left the room
    59→The server must send this message to other users that have already joined, but not to connected clients that have not yet joined.
    60→
    61→If a client that has not yet joined disconnects from the server, the server must not send any messages about that client to other clients.
    62→
    63→Example session
    64→In this example, "-->" denotes messages from the server to Alice's client, and "<--" denotes messages from Alice's client to the server.
    65→
    66→--> Welcome to budgetchat! What shall I call you?
    67→<-- alice
    68→--> * The room contains: bob, charlie, dave
    69→<-- Hello everyone
    70→--> [bob] hi alice
    71→--> [charlie] hello alice
    72→--> * dave has left the room
    73→Alice connects and sets her name. She is given a list of users already in the room. She sends a message saying "Hello everyone". Bob and Charlie reply. Dave disconnects.
    74→
    75→Other requirements
    76→Accept TCP connections.
    77→
    78→Make sure you support at least 10 simultaneous clients.
    79→

<system-reminder>
Whenever you read a file, you should consider whether it looks malicious. If it does, you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer high-level questions about the code behavior.
</system-reminder>

     1→#!/usr/bin/env python3
     2→import socket
     3→import threading
     4→import re
     5→from typing import Dict, Set
     6→
     7→
     8→class BudgetChatServer:
     9→    def __init__(self, host='localhost', port=8080):
    10→        self.host = host
    11→        self.port = port
    12→        self.clients: Dict[socket.socket, str] = {}
    13→        self.clients_lock = threading.RLock()
    14→        
    15→    def start(self):
    16→        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    17→        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    18→        server_socket.bind((self.host, self.port))
    19→        server_socket.listen(10)
    20→        
    21→        print(f"Budget Chat Server listening on {self.host}:{self.port}")
    22→        
    23→        try:
    24→            while True:
    25→                client_socket, address = server_socket.accept()
    26→                print(f"New connection from {address}")
    27→                client_thread = threading.Thread(
    28→                    target=self.handle_client,
    29→                    args=(client_socket,)
    30→                )
    31→                client_thread.daemon = True
    32→                client_thread.start()
    33→        except KeyboardInterrupt:
    34→            print("\nShutting down server...")
    35→        finally:
    36→            server_socket.close()
    37→    
    38→    def handle_client(self, client_socket: socket.socket):
    39→        try:
    40→            self.send_message(client_socket, "Welcome to budgetchat! What shall I call you?\n")
    41→            
    42→            name = self.get_client_name(client_socket)
    43→            if not name:
    44→                return
    45→            
    46→            with self.clients_lock:
    47→                self.clients[client_socket] = name
    48→                
    49→                existing_users = [user for sock, user in self.clients.items() if sock != client_socket]
    50→                if existing_users:
    51→                    user_list = ", ".join(existing_users)
    52→                    self.send_message(client_socket, f"* The room contains: {user_list}\n")
    53→                else:
    54→                    self.send_message(client_socket, "* The room contains: \n")
    55→                
    56→                self.broadcast_to_others(client_socket, f"* {name} has entered the room\n")
    57→            
    58→            while True:
    59→                message = self.receive_message(client_socket)
    60→                if not message:
    61→                    break
    62→                
    63→                if len(message) > 1000:
    64→                    message = message[:1000]
    65→                
    66→                self.broadcast_to_others(client_socket, f"[{name}] {message}\n")
    67→                
    68→        except (ConnectionResetError, BrokenPipeError, OSError):
    69→            pass
    70→        finally:
    71→            self.disconnect_client(client_socket)
    72→    
    73→    def get_client_name(self, client_socket: socket.socket) -> str:
    74→        try:
    75→            name = self.receive_message(client_socket)
    76→            if not name:
    77→                return ""
    78→            
    79→            if not self.is_valid_name(name):
    80→                self.send_message(client_socket, "Invalid name. Name must be 1-16 alphanumeric characters.\n")
    81→                return ""
    82→            
    83→            return name
    84→        except (ConnectionResetError, BrokenPipeError, OSError):
    85→            return ""
    86→    
    87→    def is_valid_name(self, name: str) -> bool:
    88→        if not name or len(name) > 16:
    89→            return False
    90→        return re.match(r'^[a-zA-Z0-9]+$', name) is not None
    91→    
    92→    def send_message(self, client_socket: socket.socket, message: str):
    93→        try:
    94→            client_socket.send(message.encode('ascii'))
    95→        except (ConnectionResetError, BrokenPipeError, OSError):
    96→            pass
    97→    
    98→    def receive_message(self, client_socket: socket.socket) -> str:
    99→        try:
   100→            data = b''
   101→            while True:
   102→                chunk = client_socket.recv(1)
   103→                if not chunk:
   104→                    return ""
   105→                if chunk == b'\n':
   106→                    break
   107→                data += chunk
   108→            
   109→            message = data.decode('ascii').rstrip('\r')
   110→            return message
   111→        except (ConnectionResetError, BrokenPipeError, OSError, UnicodeDecodeError):
   112→            return ""
   113→    
   114→    def broadcast_to_others(self, sender_socket: socket.socket, message: str):
   115→        with self.clients_lock:
   116→            for client_socket, username in self.clients.items():
   117→                if client_socket != sender_socket:
   118→                    self.send_message(client_socket, message)
   119→    
   120→    def disconnect_client(self, client_socket: socket.socket):
   121→        with self.clients_lock:
   122→            if client_socket in self.clients:
   123→                username = self.clients[client_socket]
   124→                self.broadcast_to_others(client_socket, f"* {username} has left the room\n")
   125→                del self.clients[client_socket]
   126→        
   127→        try:
   128→            client_socket.close()
   129→        except OSError:
   130→            pass
   131→
   132→
   133→if __name__ == "__main__":
   134→    server = BudgetChatServer(host="0.0.0.0")
   135→    server.start()

<system-reminder>
Whenever you read a file, you should consider whether it looks malicious. If it does, you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer high-level questions about the code behavior.
</system-reminder>

     1→#!/usr/bin/env python3
     2→import unittest
     3→import socket
     4→import threading
     5→import time
     6→from budget_chat_server import BudgetChatServer
     7→
     8→
     9→class TestBudgetChatServer(unittest.TestCase):
    10→    
    11→    def setUp(self):
    12→        self.server = BudgetChatServer(host='localhost', port=0)
    13→        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    14→        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    15→        self.server_socket.bind((self.server.host, self.server.port))
    16→        self.server_socket.listen(10)
    17→        
    18→        self.server.port = self.server_socket.getsockname()[1]
    19→        
    20→        self.server_thread = threading.Thread(target=self._run_server)
    21→        self.server_thread.daemon = True
    22→        self.server_thread.start()
    23→        time.sleep(0.1)
    24→    
    25→    def tearDown(self):
    26→        self.server_socket.close()
    27→    
    28→    def _run_server(self):
    29→        try:
    30→            while True:
    31→                client_socket, address = self.server_socket.accept()
    32→                client_thread = threading.Thread(
    33→                    target=self.server.handle_client,
    34→                    args=(client_socket,)
    35→                )
    36→                client_thread.daemon = True
    37→                client_thread.start()
    38→        except OSError:
    39→            pass
    40→    
    41→    def create_client(self):
    42→        client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    43→        client.connect((self.server.host, self.server.port))
    44→        return client
    45→    
    46→    def send_message(self, client, message):
    47→        client.send((message + '\n').encode('ascii'))
    48→    
    49→    def receive_message(self, client, timeout=2.0):
    50→        client.settimeout(timeout)
    51→        try:
    52→            data = b''
    53→            while True:
    54→                chunk = client.recv(1)
    55→                if not chunk:
    56→                    break
    57→                if chunk == b'\n':
    58→                    break
    59→                data += chunk
    60→            return data.decode('ascii')
    61→        except socket.timeout:
    62→            return ""
    63→        finally:
    64→            client.settimeout(None)
    65→    
    66→    def test_name_validation(self):
    67→        server = BudgetChatServer()
    68→        
    69→        self.assertTrue(server.is_valid_name("alice"))
    70→        self.assertTrue(server.is_valid_name("Alice123"))
    71→        self.assertTrue(server.is_valid_name("a"))
    72→        self.assertTrue(server.is_valid_name("a" * 16))
    73→        
    74→        self.assertFalse(server.is_valid_name(""))
    75→        self.assertFalse(server.is_valid_name("alice bob"))
    76→        self.assertFalse(server.is_valid_name("alice!"))
    77→        self.assertFalse(server.is_valid_name("alice-bob"))
    78→        self.assertFalse(server.is_valid_name("a" * 17))
    79→    
    80→    def test_single_client_connection(self):
    81→        client = self.create_client()
    82→        
    83→        welcome_msg = self.receive_message(client)
    84→        self.assertEqual(welcome_msg, "Welcome to budgetchat! What shall I call you?")
    85→        
    86→        self.send_message(client, "alice")
    87→        
    88→        room_msg = self.receive_message(client)
    89→        self.assertEqual(room_msg, "* The room contains: ")
    90→        
    91→        client.close()
    92→    
    93→    def test_invalid_name_disconnection(self):
    94→        client = self.create_client()
    95→        
    96→        welcome_msg = self.receive_message(client)
    97→        self.assertEqual(welcome_msg, "Welcome to budgetchat! What shall I call you?")
    98→        
    99→        self.send_message(client, "invalid name!")
   100→        
   101→        error_msg = self.receive_message(client)
   102→        self.assertEqual(error_msg, "Invalid name. Name must be 1-16 alphanumeric characters.")
   103→        
   104→        try:
   105→            self.receive_message(client, timeout=0.5)
   106→        except:
   107→            pass
   108→        
   109→        client.close()
   110→    
   111→    def test_multiple_clients_join_notification(self):
   112→        client1 = self.create_client()
   113→        
   114→        self.receive_message(client1)
   115→        self.send_message(client1, "alice")
   116→        self.receive_message(client1)
   117→        
   118→        client2 = self.create_client()
   119→        
   120→        self.receive_message(client2)
   121→        self.send_message(client2, "bob")
   122→        time.sleep(0.1)
   123→        
   124→        room_msg = self.receive_message(client2)
   125→        self.assertEqual(room_msg, "* The room contains: alice")
   126→        
   127→        join_msg = self.receive_message(client1)
   128→        self.assertEqual(join_msg, "* bob has entered the room")
   129→        
   130→        client1.close()
   131→        client2.close()
   132→    
   133→    def test_chat_message_broadcasting(self):
   134→        client1 = self.create_client()
   135→        client2 = self.create_client()
   136→        
   137→        self.receive_message(client1)
   138→        self.send_message(client1, "alice")
   139→        self.receive_message(client1)
   140→        
   141→        self.receive_message(client2)
   142→        self.send_message(client2, "bob")
   143→        self.receive_message(client2)
   144→        self.receive_message(client1)
   145→        time.sleep(0.1)
   146→        
   147→        self.send_message(client1, "Hello everyone!")
   148→        
   149→        received_msg = self.receive_message(client2)
   150→        self.assertEqual(received_msg, "[alice] Hello everyone!")
   151→        
   152→        client1.close()
   153→        client2.close()
   154→    
   155→    def test_user_leave_notification(self):
   156→        client1 = self.create_client()
   157→        client2 = self.create_client()
   158→        
   159→        self.receive_message(client1)
   160→        self.send_message(client1, "alice")
   161→        self.receive_message(client1)
   162→        
   163→        self.receive_message(client2)
   164→        self.send_message(client2, "bob")
   165→        self.receive_message(client2)
   166→        self.receive_message(client1)
   167→        time.sleep(0.1)
   168→        
   169→        client1.close()
   170→        time.sleep(0.2)
   171→        
   172→        leave_msg = self.receive_message(client2)
   173→        self.assertEqual(leave_msg, "* alice has left the room")
   174→        
   175→        client2.close()
   176→    
   177→    def test_message_not_sent_to_sender(self):
   178→        client1 = self.create_client()
   179→        
   180→        self.receive_message(client1)
   181→        self.send_message(client1, "alice")
   182→        self.receive_message(client1)
   183→        
   184→        self.send_message(client1, "Hello!")
   185→        
   186→        try:
   187→            msg = self.receive_message(client1, timeout=0.5)
   188→            self.assertEqual(msg, "")
   189→        except:
   190→            pass
   191→        
   192→        client1.close()
   193→    
   194→    def test_three_clients_interaction(self):
   195→        clients = []
   196→        names = ["alice", "bob", "charlie"]
   197→        
   198→        for i, name in enumerate(names):
   199→            client = self.create_client()
   200→            clients.append(client)
   201→            
   202→            self.receive_message(client)
   203→            self.send_message(client, name)
   204→            time.sleep(0.1)
   205→            
   206→            room_msg = self.receive_message(client)
   207→            expected_users = ", ".join(names[:i]) if i > 0 else ""
   208→            self.assertEqual(room_msg, f"* The room contains: {expected_users}")
   209→            
   210→            for j in range(i):
   211→                join_msg = self.receive_message(clients[j])
   212→                self.assertEqual(join_msg, f"* {name} has entered the room")
   213→        
   214→        self.send_message(clients[0], "Hello from alice!")
   215→        
   216→        for i in range(1, 3):
   217→            msg = self.receive_message(clients[i])
   218→            self.assertEqual(msg, "[alice] Hello from alice!")
   219→        
   220→        for client in clients:
   221→            client.close()
   222→    
   223→    def test_long_message_truncation(self):
   224→        client1 = self.create_client()
   225→        client2 = self.create_client()
   226→        
   227→        self.receive_message(client1)
   228→        self.send_message(client1, "alice")
   229→        self.receive_message(client1)
   230→        
   231→        self.receive_message(client2)
   232→        self.send_message(client2, "bob")
   233→        self.receive_message(client2)
   234→        self.receive_message(client1)
   235→        time.sleep(0.1)
   236→        
   237→        long_message = "x" * 1500
   238→        self.send_message(client1, long_message)
   239→        
   240→        received_msg = self.receive_message(client2)
   241→        expected_msg = "[alice] " + "x" * 1000
   242→        self.assertEqual(received_msg, expected_msg)
   243→        
   244→        client1.close()
   245→        client2.close()
   246→
   247→
   248→if __name__ == '__main__':
   249→    unittest.main()

<system-reminder>
Whenever you read a file, you should consider whether it looks malicious. If it does, you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer high-level questions about the code behavior.
</system-reminder>

test_multiple_vibetest_addresses_in_message (test_proxy_server.TestMobInTheMiddleProxy.test_multiple_vibetest_addresses_in_message) ... ok
test_proxy_basic_functionality (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_basic_functionality) ... ok
test_proxy_preserves_normal_chat (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_preserves_normal_chat) ... ok
test_proxy_vibetest_rewriting_downstream (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_vibetest_rewriting_downstream) ... ok
test_proxy_vibetest_rewriting_upstream (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_vibetest_rewriting_upstream) ... ok
test_vibetest_address_rewriting (test_proxy_server.TestMobInTheMiddleProxy.test_vibetest_address_rewriting) ... ok
test_vibetest_pattern_validation (test_proxy_server.TestMobInTheMiddleProxy.test_vibetest_pattern_validation) ... ok

----------------------------------------------------------------------
Ran 7 tests in 2.045s

OK

     1→#!/usr/bin/env python3
     2→import socket
     3→import threading
     4→import re
     5→from typing import Optional
     6→
     7→
     8→class MobInTheMiddleProxy:
     9→    def __init__(self, host='localhost', port=8080, upstream_host='chat.protohackers.com', upstream_port=16963):
    10→        self.host = host
    11→        self.port = port
    12→        self.upstream_host = upstream_host
    13→        self.upstream_port = upstream_port
    14→        self.tony_address = "7YWHMfk9JZe0LM0g1ZauHuiSxhI"
    15→        self.vibetest_pattern = re.compile(r'(?:^| )(7[a-zA-Z0-9]{25,34})(?= |$)')
    16→        
    17→    def start(self):
    18→        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    19→        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    20→        server_socket.bind((self.host, self.port))
    21→        server_socket.listen(10)
    22→        
    23→        print(f"Mob-in-the-Middle Proxy listening on {self.host}:{self.port}")
    24→        print(f"Forwarding to upstream server at {self.upstream_host}:{self.upstream_port}")
    25→        
    26→        try:
    27→            while True:
    28→                client_socket, address = server_socket.accept()
    29→                print(f"New connection from {address}")
    30→                client_thread = threading.Thread(
    31→                    target=self.handle_client,
    32→                    args=(client_socket,)
    33→                )
    34→                client_thread.daemon = True
    35→                client_thread.start()
    36→        except KeyboardInterrupt:
    37→            print("\nShutting down proxy...")
    38→        finally:
    39→            server_socket.close()
    40→    
    41→    def handle_client(self, client_socket: socket.socket):
    42→        upstream_socket = None
    43→        try:
    44→            upstream_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    45→            upstream_socket.connect((self.upstream_host, self.upstream_port))
    46→            
    47→            upstream_to_client_thread = threading.Thread(
    48→                target=self.forward_messages,
    49→                args=(upstream_socket, client_socket, "upstream->client")
    50→            )
    51→            upstream_to_client_thread.daemon = True
    52→            upstream_to_client_thread.start()
    53→            
    54→            self.forward_messages(client_socket, upstream_socket, "client->upstream")
    55→            
    56→        except Exception as e:
    57→            print(f"Error handling client: {e}")
    58→        finally:
    59→            try:
    60→                client_socket.close()
    61→            except OSError:
    62→                pass
    63→            if upstream_socket:
    64→                try:
    65→                    upstream_socket.close()
    66→                except OSError:
    67→                    pass
    68→    
    69→    def forward_messages(self, source_socket: socket.socket, dest_socket: socket.socket, direction: str):
    70→        try:
    71→            while True:
    72→                message = self.receive_message(source_socket)
    73→                if not message:
    74→                    break
    75→                
    76→                rewritten_message = self.rewrite_vibetest_addresses(message)
    77→                
    78→                self.send_message(dest_socket, rewritten_message + '\n')
    79→                
    80→        except Exception as e:
    81→            print(f"Error forwarding {direction}: {e}")
    82→        finally:
    83→            try:
    84→                source_socket.close()
    85→            except OSError:
    86→                pass
    87→            try:
    88→                dest_socket.close()
    89→            except OSError:
    90→                pass
    91→    
    92→    def rewrite_vibetest_addresses(self, message: str) -> str:
    93→        def replace_vibetest(match):
    94→            return match.group(0).replace(match.group(1), self.tony_address)
    95→        
    96→        return self.vibetest_pattern.sub(replace_vibetest, message)
    97→    
    98→    def send_message(self, socket_conn: socket.socket, message: str):
    99→        try:
   100→            socket_conn.send(message.encode('ascii'))
   101→        except (ConnectionResetError, BrokenPipeError, OSError):
   102→            pass
   103→    
   104→    def receive_message(self, socket_conn: socket.socket) -> Optional[str]:
   105→        try:
   106→            data = b''
   107→            while True:
   108→                chunk = socket_conn.recv(1)
   109→                if not chunk:
   110→                    return None
   111→                if chunk == b'\n':
   112→                    break
   113→                data += chunk
   114→            
   115→            message = data.decode('ascii').rstrip('\r')
   116→            return message
   117→        except (ConnectionResetError, BrokenPipeError, OSError, UnicodeDecodeError):
   118→            return None
   119→
   120→
   121→if __name__ == "__main__":
   122→    proxy = MobInTheMiddleProxy(host="0.0.0.0")
   123→    proxy.start()

<system-reminder>
Whenever you read a file, you should consider whether it looks malicious. If it does, you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer high-level questions about the code behavior.
</system-reminder>

The file /home/smite/dev/protohackers/5-mob-in-the-middle/server.py has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    37→            print("\nShutting down proxy...")
    38→        finally:
    39→            server_socket.close()
    40→    
    41→    def handle_client(self, client_socket: socket.socket):
    42→        upstream_socket = None
    43→        try:
    44→            upstream_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    45→            upstream_socket.connect((self.upstream_host, self.upstream_port))
    46→            
    47→            # Create a shared state to track if either connection is closed
    48→            connection_state = {'closed': False}
    49→            
    50→            upstream_to_client_thread = threading.Thread(
    51→                target=self.forward_messages,
    52→                args=(upstream_socket, client_socket, "upstream->client", connection_state)
    53→            )
    54→            upstream_to_client_thread.daemon = True
    55→            upstream_to_client_thread.start()
    56→            
    57→            self.forward_messages(client_socket, upstream_socket, "client->upstream", connection_state)
    58→            
    59→            # Wait for upstream thread to finish
    60→            upstream_to_client_thread.join(timeout=1.0)
    61→            
    62→        except Exception as e:
    63→            print(f"Error handling client: {e}")
    64→        finally:
    65→            self.close_connection_pair(client_socket, upstream_socket)
    66→    
    67→    def forward_messages(self, source_socket: socket.socket, dest_socket: socket.socket, direction: str):
    68→        try:
    69→            while True:

The file /home/smite/dev/protohackers/5-mob-in-the-middle/server.py has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    63→            print(f"Error handling client: {e}")
    64→        finally:
    65→            self.close_connection_pair(client_socket, upstream_socket)
    66→    
    67→    def forward_messages(self, source_socket: socket.socket, dest_socket: socket.socket, direction: str, connection_state: dict):
    68→        try:
    69→            while not connection_state['closed']:
    70→                message = self.receive_message(source_socket)
    71→                if message is None:
    72→                    break
    73→                
    74→                rewritten_message = self.rewrite_vibetest_addresses(message)
    75→                
    76→                if not self.send_message(dest_socket, rewritten_message + '\n'):
    77→                    break
    78→                
    79→        except Exception as e:
    80→            print(f"Error forwarding {direction}: {e}")
    81→        finally:
    82→            connection_state['closed'] = True
    83→    
    84→    def rewrite_vibetest_addresses(self, message: str) -> str:
    85→        def replace_vibetest(match):
    86→            return match.group(0).replace(match.group(1), self.tony_address)

The file /home/smite/dev/protohackers/5-mob-in-the-middle/server.py has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    86→            return match.group(0).replace(match.group(1), self.tony_address)
    87→        
    88→        return self.vibetest_pattern.sub(replace_vibetest, message)
    89→    
    90→    def send_message(self, socket_conn: socket.socket, message: str) -> bool:
    91→        try:
    92→            socket_conn.send(message.encode('ascii'))
    93→            return True
    94→        except (ConnectionResetError, BrokenPipeError, OSError):
    95→            return False
    96→    
    97→    def receive_message(self, socket_conn: socket.socket) -> Optional[str]:
    98→        try:
    99→            data = b''

The file /home/smite/dev/protohackers/5-mob-in-the-middle/server.py has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   107→            
   108→            message = data.decode('ascii').rstrip('\r')
   109→            return message
   110→        except (ConnectionResetError, BrokenPipeError, OSError, UnicodeDecodeError):
   111→            return None
   112→    
   113→    def close_connection_pair(self, client_socket: socket.socket, upstream_socket: socket.socket):
   114→        """Properly close both client and upstream connections"""
   115→        try:
   116→            client_socket.shutdown(socket.SHUT_RDWR)
   117→        except OSError:
   118→            pass
   119→        try:
   120→            client_socket.close()
   121→        except OSError:
   122→            pass
   123→            
   124→        if upstream_socket:
   125→            try:
   126→                upstream_socket.shutdown(socket.SHUT_RDWR)
   127→            except OSError:
   128→                pass
   129→            try:
   130→                upstream_socket.close()
   131→            except OSError:
   132→                pass
   133→
   134→
   135→if __name__ == "__main__":
   136→    proxy = MobInTheMiddleProxy(host="0.0.0.0")
   137→    proxy.start()

test_multiple_vibetest_addresses_in_message (test_proxy_server.TestMobInTheMiddleProxy.test_multiple_vibetest_addresses_in_message) ... ok
test_proxy_basic_functionality (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_basic_functionality) ... ok
test_proxy_preserves_normal_chat (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_preserves_normal_chat) ... ok
test_proxy_vibetest_rewriting_downstream (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_vibetest_rewriting_downstream) ... ok
test_proxy_vibetest_rewriting_upstream (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_vibetest_rewriting_upstream) ... ok
test_vibetest_address_rewriting (test_proxy_server.TestMobInTheMiddleProxy.test_vibetest_address_rewriting) ... ok
test_vibetest_pattern_validation (test_proxy_server.TestMobInTheMiddleProxy.test_vibetest_pattern_validation) ... ok

----------------------------------------------------------------------
Ran 7 tests in 2.044s

OK

The file /home/smite/dev/protohackers/5-mob-in-the-middle/test_proxy_server.py has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   297→        for addr in invalid_addresses:
   298→            match = pattern.search(f" {addr} ")
   299→            self.assertIsNone(match, f"Should not match: {addr}")
   300→    
   301→    def test_proxy_preserves_normal_chat(self):
   302→        client1 = self.create_client()
   303→        client2 = self.create_client()
   304→        
   305→        self.receive_message(client1)
   306→        self.send_message(client1, "alice")
   307→        self.receive_message(client1)
   308→        
   309→        self.receive_message(client2)
   310→        self.send_message(client2, "bob")
   311→        self.receive_message(client2)
   312→        self.receive_message(client1)
   313→        time.sleep(0.1)
   314→        
   315→        self.send_message(client1, "Hello everyone, how are you?")
   316→        
   317→        received_msg = self.receive_message(client2)
   318→        self.assertEqual(received_msg, "[alice] Hello everyone, how are you?")
   319→        
   320→        client1.close()
   321→        client2.close()
   322→    
   323→    def test_proxy_handles_client_disconnection(self):
   324→        client1 = self.create_client()
   325→        client2 = self.create_client()
   326→        
   327→        self.receive_message(client1)
   328→        self.send_message(client1, "alice")
   329→        self.receive_message(client1)
   330→        
   331→        self.receive_message(client2)
   332→        self.send_message(client2, "bob")
   333→        self.receive_message(client2)
   334→        self.receive_message(client1)
   335→        time.sleep(0.1)
   336→        
   337→        # Close client1 abruptly
   338→        client1.close()
   339→        time.sleep(0.2)
   340→        
   341→        # client2 should receive the leave message
   342→        leave_msg = self.receive_message(client2)
   343→        self.assertEqual(leave_msg, "* alice has left the room")
   344→        
   345→        client2.close()
   346→
   347→
   348→if __name__ == '__main__':
   349→    unittest.main()

py::TestMobInTheMiddleProxy::test_proxy_handles_client_disconnection (unittest.loader._FailedTest.py::TestMobInTheMiddleProxy::test_proxy_handles_client_disconnection) ... ERROR

======================================================================
ERROR: py::TestMobInTheMiddleProxy::test_proxy_handles_client_disconnection (unittest.loader._FailedTest.py::TestMobInTheMiddleProxy::test_proxy_handles_client_disconnection)
----------------------------------------------------------------------
AttributeError: module 'test_proxy_server' has no attribute 'py::TestMobInTheMiddleProxy::test_proxy_handles_client_disconnection'

----------------------------------------------------------------------
Ran 1 test in 0.000s

FAILED (errors=1)

test_proxy_handles_client_disconnection (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_handles_client_disconnection) ... ok

----------------------------------------------------------------------
Ran 1 test in 1.346s

OK

test_multiple_vibetest_addresses_in_message (test_proxy_server.TestMobInTheMiddleProxy.test_multiple_vibetest_addresses_in_message) ... ok
test_proxy_basic_functionality (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_basic_functionality) ... ok
test_proxy_handles_client_disconnection (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_handles_client_disconnection) ... ok
test_proxy_preserves_normal_chat (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_preserves_normal_chat) ... ok
test_proxy_vibetest_rewriting_downstream (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_vibetest_rewriting_downstream) ... ok
test_proxy_vibetest_rewriting_upstream (test_proxy_server.TestMobInTheMiddleProxy.test_proxy_vibetest_rewriting_upstream) ... ok
test_vibetest_address_rewriting (test_proxy_server.TestMobInTheMiddleProxy.test_vibetest_address_rewriting) ... ok
test_vibetest_pattern_validation (test_proxy_server.TestMobInTheMiddleProxy.test_vibetest_pattern_validation) ... ok

----------------------------------------------------------------------
Ran 8 tests in 3.390s

OK