tests/tinyproxy.py
changeset 43076 2372284d9457
parent 41709 97e2442a4595
child 45830 c102b704edb5
equal deleted inserted replaced
43075:57875cf423c9 43076:2372284d9457
    32 if os.environ.get('HGIPV6', '0') == '1':
    32 if os.environ.get('HGIPV6', '0') == '1':
    33     family = socket.AF_INET6
    33     family = socket.AF_INET6
    34 else:
    34 else:
    35     family = socket.AF_INET
    35     family = socket.AF_INET
    36 
    36 
    37 class ProxyHandler (httpserver.basehttprequesthandler):
    37 
       
    38 class ProxyHandler(httpserver.basehttprequesthandler):
    38     __base = httpserver.basehttprequesthandler
    39     __base = httpserver.basehttprequesthandler
    39     __base_handle = __base.handle
    40     __base_handle = __base.handle
    40 
    41 
    41     server_version = "TinyHTTPProxy/" + __version__
    42     server_version = "TinyHTTPProxy/" + __version__
    42     rbufsize = 0                        # self.rfile Be unbuffered
    43     rbufsize = 0  # self.rfile Be unbuffered
    43 
    44 
    44     def handle(self):
    45     def handle(self):
    45         (ip, port) = self.client_address
    46         (ip, port) = self.client_address
    46         allowed = getattr(self, 'allowed_clients', None)
    47         allowed = getattr(self, 'allowed_clients', None)
    47         if allowed is not None and ip not in allowed:
    48         if allowed is not None and ip not in allowed:
    51         else:
    52         else:
    52             self.__base_handle()
    53             self.__base_handle()
    53 
    54 
    54     def log_request(self, code='-', size='-'):
    55     def log_request(self, code='-', size='-'):
    55         xheaders = [h for h in self.headers.items() if h[0].startswith('x-')]
    56         xheaders = [h for h in self.headers.items() if h[0].startswith('x-')]
    56         self.log_message('"%s" %s %s%s',
    57         self.log_message(
    57                          self.requestline, str(code), str(size),
    58             '"%s" %s %s%s',
    58                          ''.join([' %s:%s' % h for h in sorted(xheaders)]))
    59             self.requestline,
       
    60             str(code),
       
    61             str(size),
       
    62             ''.join([' %s:%s' % h for h in sorted(xheaders)]),
       
    63         )
    59         # Flush for Windows, so output isn't lost on TerminateProcess()
    64         # Flush for Windows, so output isn't lost on TerminateProcess()
    60         sys.stdout.flush()
    65         sys.stdout.flush()
    61         sys.stderr.flush()
    66         sys.stderr.flush()
    62 
    67 
    63     def _connect_to(self, netloc, soc):
    68     def _connect_to(self, netloc, soc):
    64         i = netloc.find(':')
    69         i = netloc.find(':')
    65         if i >= 0:
    70         if i >= 0:
    66             host_port = netloc[:i], int(netloc[i + 1:])
    71             host_port = netloc[:i], int(netloc[i + 1 :])
    67         else:
    72         else:
    68             host_port = netloc, 80
    73             host_port = netloc, 80
    69         print("\t" "connect to %s:%d" % host_port)
    74         print("\t" "connect to %s:%d" % host_port)
    70         try: soc.connect(host_port)
    75         try:
       
    76             soc.connect(host_port)
    71         except socket.error as arg:
    77         except socket.error as arg:
    72             try: msg = arg[1]
    78             try:
    73             except (IndexError, TypeError): msg = arg
    79                 msg = arg[1]
       
    80             except (IndexError, TypeError):
       
    81                 msg = arg
    74             self.send_error(404, msg)
    82             self.send_error(404, msg)
    75             return 0
    83             return 0
    76         return 1
    84         return 1
    77 
    85 
    78     def do_CONNECT(self):
    86     def do_CONNECT(self):
    79         soc = socket.socket(family, socket.SOCK_STREAM)
    87         soc = socket.socket(family, socket.SOCK_STREAM)
    80         try:
    88         try:
    81             if self._connect_to(self.path, soc):
    89             if self._connect_to(self.path, soc):
    82                 self.log_request(200)
    90                 self.log_request(200)
    83                 self.wfile.write(pycompat.bytestr(self.protocol_version) +
    91                 self.wfile.write(
    84                                  b" 200 Connection established\r\n")
    92                     pycompat.bytestr(self.protocol_version)
    85                 self.wfile.write(b"Proxy-agent: %s\r\n" %
    93                     + b" 200 Connection established\r\n"
    86                                  pycompat.bytestr(self.version_string()))
    94                 )
       
    95                 self.wfile.write(
       
    96                     b"Proxy-agent: %s\r\n"
       
    97                     % pycompat.bytestr(self.version_string())
       
    98                 )
    87                 self.wfile.write(b"\r\n")
    99                 self.wfile.write(b"\r\n")
    88                 self._read_write(soc, 300)
   100                 self._read_write(soc, 300)
    89         finally:
   101         finally:
    90             print("\t" "bye")
   102             print("\t" "bye")
    91             soc.close()
   103             soc.close()
    92             self.connection.close()
   104             self.connection.close()
    93 
   105 
    94     def do_GET(self):
   106     def do_GET(self):
    95         (scm, netloc, path, params, query, fragment) = urlreq.urlparse(
   107         (scm, netloc, path, params, query, fragment) = urlreq.urlparse(
    96             self.path, 'http')
   108             self.path, 'http'
       
   109         )
    97         if scm != 'http' or fragment or not netloc:
   110         if scm != 'http' or fragment or not netloc:
    98             self.send_error(400, "bad url %s" % self.path)
   111             self.send_error(400, "bad url %s" % self.path)
    99             return
   112             return
   100         soc = socket.socket(family, socket.SOCK_STREAM)
   113         soc = socket.socket(family, socket.SOCK_STREAM)
   101         try:
   114         try:
   102             if self._connect_to(netloc, soc):
   115             if self._connect_to(netloc, soc):
   103                 self.log_request()
   116                 self.log_request()
   104                 url = urlreq.urlunparse(('', '', path, params, query, ''))
   117                 url = urlreq.urlunparse(('', '', path, params, query, ''))
   105                 soc.send(b"%s %s %s\r\n" % (
   118                 soc.send(
   106                     pycompat.bytestr(self.command),
   119                     b"%s %s %s\r\n"
   107                     pycompat.bytestr(url),
   120                     % (
   108                     pycompat.bytestr(self.request_version)))
   121                         pycompat.bytestr(self.command),
       
   122                         pycompat.bytestr(url),
       
   123                         pycompat.bytestr(self.request_version),
       
   124                     )
       
   125                 )
   109                 self.headers['Connection'] = 'close'
   126                 self.headers['Connection'] = 'close'
   110                 del self.headers['Proxy-Connection']
   127                 del self.headers['Proxy-Connection']
   111                 for key, val in self.headers.items():
   128                 for key, val in self.headers.items():
   112                     soc.send(b"%s: %s\r\n" % (pycompat.bytestr(key),
   129                     soc.send(
   113                                               pycompat.bytestr(val)))
   130                         b"%s: %s\r\n"
       
   131                         % (pycompat.bytestr(key), pycompat.bytestr(val))
       
   132                     )
   114                 soc.send(b"\r\n")
   133                 soc.send(b"\r\n")
   115                 self._read_write(soc)
   134                 self._read_write(soc)
   116         finally:
   135         finally:
   117             print("\t" "bye")
   136             print("\t" "bye")
   118             soc.close()
   137             soc.close()
   145             if count == max_idling:
   164             if count == max_idling:
   146                 break
   165                 break
   147 
   166 
   148     do_HEAD = do_GET
   167     do_HEAD = do_GET
   149     do_POST = do_GET
   168     do_POST = do_GET
   150     do_PUT  = do_GET
   169     do_PUT = do_GET
   151     do_DELETE = do_GET
   170     do_DELETE = do_GET
   152 
   171 
   153 class ThreadingHTTPServer (socketserver.ThreadingMixIn,
   172 
   154                            httpserver.httpserver):
   173 class ThreadingHTTPServer(socketserver.ThreadingMixIn, httpserver.httpserver):
   155     def __init__(self, *args, **kwargs):
   174     def __init__(self, *args, **kwargs):
   156         httpserver.httpserver.__init__(self, *args, **kwargs)
   175         httpserver.httpserver.__init__(self, *args, **kwargs)
   157         a = open("proxy.pid", "w")
   176         a = open("proxy.pid", "w")
   158         a.write(str(os.getpid()) + "\n")
   177         a.write(str(os.getpid()) + "\n")
   159         a.close()
   178         a.close()
       
   179 
   160 
   180 
   161 def runserver(port=8000, bind=""):
   181 def runserver(port=8000, bind=""):
   162     server_address = (bind, port)
   182     server_address = (bind, port)
   163     ProxyHandler.protocol_version = "HTTP/1.0"
   183     ProxyHandler.protocol_version = "HTTP/1.0"
   164     httpd = ThreadingHTTPServer(server_address, ProxyHandler)
   184     httpd = ThreadingHTTPServer(server_address, ProxyHandler)
   168         httpd.serve_forever()
   188         httpd.serve_forever()
   169     except KeyboardInterrupt:
   189     except KeyboardInterrupt:
   170         print("\nKeyboard interrupt received, exiting.")
   190         print("\nKeyboard interrupt received, exiting.")
   171         httpd.server_close()
   191         httpd.server_close()
   172         sys.exit(0)
   192         sys.exit(0)
       
   193 
   173 
   194 
   174 if __name__ == '__main__':
   195 if __name__ == '__main__':
   175     argv = sys.argv
   196     argv = sys.argv
   176     if argv[1:] and argv[1] in ('-h', '--help'):
   197     if argv[1:] and argv[1] in ('-h', '--help'):
   177         print(argv[0], "[port [allowed_client_name ...]]")
   198         print(argv[0], "[port [allowed_client_name ...]]")
   186             del argv[2:]
   207             del argv[2:]
   187         else:
   208         else:
   188             print("Any clients will be served...")
   209             print("Any clients will be served...")
   189 
   210 
   190         parser = optparse.OptionParser()
   211         parser = optparse.OptionParser()
   191         parser.add_option('-b', '--bind', metavar='ADDRESS',
   212         parser.add_option(
   192                           help='Specify alternate bind address '
   213             '-b',
   193                                '[default: all interfaces]', default='')
   214             '--bind',
       
   215             metavar='ADDRESS',
       
   216             help='Specify alternate bind address ' '[default: all interfaces]',
       
   217             default='',
       
   218         )
   194         (options, args) = parser.parse_args()
   219         (options, args) = parser.parse_args()
   195         port = 8000
   220         port = 8000
   196         if len(args) == 1:
   221         if len(args) == 1:
   197             port = int(args[0])
   222             port = int(args[0])
   198         runserver(port, options.bind)
   223         runserver(port, options.bind)