Python Server and Client

import socket
import sys
import threading
import base64
import json
from time import sleep

class PDU:
    def __init__(self, msg_type, body):
        self.header = {"msg_type": msg_type}
        self.body = body

    def to_dict(self):
        return {"header": self.header, "body": self.body}

    @classmethod
    def from_dict(cls, data):
        return cls(msg_type=data["header"]["msg_type"], body=data["body"])

def server_program(port):
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind(("0.0.0.0", port))
    server_socket.listen(1)
    print(f"Server listening on port {port}")

    conn, addr = server_socket.accept()
    print(f"Connection established with {addr}")

    while True:
        data = conn.recv(1024)
        if not data:
            break
        try:
            pdu = PDU.from_dict(json.loads(data.decode()))
            if pdu.header["msg_type"] == "msg":
                message = base64.b64decode(pdu.body).decode()
                print(f"Received: {message}")

                # Send acknowledgment
                ack_pdu = PDU("ack", None)
                conn.send(json.dumps(ack_pdu.to_dict()).encode())
        except Exception as e:
            print(f"Error processing message: {e}")
            break

    conn.close()
    server_socket.close()
    print("Server closed.")

def client_program(ip, port):
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect((ip, port))
    print(f"Connected to server at {ip}:{port}")

    retries = 3
    while True:
        message = input("Enter message (type 'close' to exit): ")
        if message.lower() == "close":
            break

        pdu = PDU("msg", base64.b64encode(message.encode()).decode())
        for attempt in range(retries):
            client_socket.send(json.dumps(pdu.to_dict()).encode())

            client_socket.settimeout(5)
            try:
                ack_data = client_socket.recv(1024)
                ack_pdu = PDU.from_dict(json.loads(ack_data.decode()))
                if ack_pdu.header["msg_type"] == "ack":
                    print("Acknowledgment received.")
                    break
            except socket.timeout:
                print(f"No acknowledgment. Retrying ({attempt + 1}/{retries})...")
        else:
            print("Failed to receive acknowledgment after 3 attempts.")

    client_socket.close()
    print("Client closed.")

def main():
    role = input("Are you running the server or client? (server/client): ").strip().lower()
    if role == "server":
        port = int(input("Enter port number: "))
        server_program(port)
    elif role == "client":
        ip = input("Enter server IP address: ")
        port = int(input("Enter server port number: "))
        client_program(ip, port)
    else:
        print("Invalid role. Please restart.")

if __name__ == "__main__":
    main()