1  from argparse import ArgumentParser
 2  from socket import *
 3  from traceback import format_exc
 4  
 5  
 6  class BaseHandler:
 7      def handle(self, c):
 8          print(c)
 9  
10  
11  class BaseServer:
12      def __init__(self, bind, handler):
13          self.bind = bind
14          self.handler = handler
15  
16      def start(self):
17          pass
18  
19  
20  class TCPServer(BaseServer):
21      def __init__(self, bind, handler):
22          super().__init__(bind, handler)
23          self.s = socket(AF_INET, SOCK_STREAM)
24          self.s.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
25          self.s.bind(bind)
26          self.s.listen(1)
27  
28      def start(self):
29          try:
30              while True:
31                  c, caddr = self.s.accept()
32                  c.settimeout(10)
33                  print("Connected:", caddr)
34  
35                  try:
36                      self.handler.handle(c)
37                  except TimeoutError:
38                      print("Timeout")
39                      pass
40                  except Exception as e:
41                      print("\x1b[31m" + format_exc() + "\x1b[m")
42  
43                  print("Disconnected")
44          except KeyboardInterrupt:
45              self.s.close()
46  
47  
48  class UDPServer(BaseServer):
49      def __init__(self, bind, handler):
50          super().__init__(bind, handler)
51          self.s = socket(AF_INET, SOCK_DGRAM)
52          self.s.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
53          self.s.bind(bind)
54  
55      def start(self):
56          try:
57              while True:
58                  try:
59                      self.handler.handle(self.s)
60                  except Exception as e:
61                      print("\x1b[31m" + format_exc() + "\x1b[m")
62  
63          except KeyboardInterrupt:
64              self.s.close()
65  
66  
67  def server(mode, bind, handler):
68      if mode == "tcp":
69          server = TCPServer
70      elif mode == "udp":
71          server = UDPServer
72      else:
73          raise ValueError("Server mode can be either tcp or udp")
74  
75      print(f"Opening {mode.upper()} server at {bind[0]}:{bind[1]}")
76  
77      server(bind, handler).start()
78  
79  
80  if __name__ == '__main__':
81      parser = ArgumentParser()
82      parser.add_argument("--mode", choice=("tcp", "udp"), default="tcp")
83      parser.add_argument("--port", type=int, default=1234)
84      args = parser.parse_args()
85  
86      server(args.mode, ('127.0.0.1', args.port), BaseHandler())