Cleaned up the code a bit and made a few bug fixes

This commit is contained in:
allfro 2012-12-24 00:52:23 -05:00
parent 225b071883
commit f202fad441

View File

@ -1,11 +1,11 @@
#!/usr/bin/env python #!/usr/bin/env python
from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler
from urlparse import urlparse, urlunparse, ParseResult
from SocketServer import ThreadingMixIn from SocketServer import ThreadingMixIn
from os import path, mkdir, listdir
from httplib import HTTPResponse from httplib import HTTPResponse
from tempfile import gettempdir from tempfile import gettempdir
from urlparse import urlparse from os import path, listdir
from ssl import wrap_socket from ssl import wrap_socket
from socket import socket from socket import socket
from re import compile from re import compile
@ -129,21 +129,35 @@ class ProxyHandler(BaseHTTPRequestHandler):
r = compile(r'http://[^/]+(/?.*)(?i)') r = compile(r'http://[^/]+(/?.*)(?i)')
def __init__(self, request, client_address, server):
self.is_connect = False
BaseHTTPRequestHandler.__init__(self, request, client_address, server)
def _connect_to_host(self): def _connect_to_host(self):
# Get hostname and port to connect to # Get hostname and port to connect to
if self.is_connect: if self.is_connect:
host, port = self.path.split(':') self.hostname, self.port = self.path.split(':')
else: else:
u = urlparse(self.path) u = urlparse(self.path)
if u.scheme != 'http': if u.scheme != 'http':
raise UnsupportedSchemeException('Unknown scheme %s' % repr(u.scheme)) raise UnsupportedSchemeException('Unknown scheme %s' % repr(u.scheme))
host = u.hostname self.hostname = u.hostname
port = u.port or 80 self.port = u.port or 80
self.path = urlunparse(
ParseResult(
scheme='',
netloc='',
params=u.params,
path=u.path or '/',
query=u.query,
fragment=u.fragment
)
)
# Connect to destination # Connect to destination
self._proxy_sock = socket() self._proxy_sock = socket()
self._proxy_sock.settimeout(10) self._proxy_sock.settimeout(10)
self._proxy_sock.connect((host, int(port))) self._proxy_sock.connect((self.hostname, int(self.port)))
# Wrap socket if SSL is required # Wrap socket if SSL is required
if self.is_connect: if self.is_connect:
@ -178,8 +192,6 @@ class ProxyHandler(BaseHTTPRequestHandler):
def do_COMMAND(self): def do_COMMAND(self):
# Is this an SSL tunnel? # Is this an SSL tunnel?
path = self.path or '/'
if not self.is_connect: if not self.is_connect:
try: try:
# Connect to destination # Connect to destination
@ -188,10 +200,9 @@ class ProxyHandler(BaseHTTPRequestHandler):
self.send_error(500, str(e)) self.send_error(500, str(e))
return return
# Extract path # Extract path
path = self.r.search(path).groups()[0] or '/'
# Build request # Build request
req = '%s %s %s\r\n' % (self.command, path, self.request_version) req = '%s %s %s\r\n' % (self.command, self.path, self.request_version)
# Add headers to the request # Add headers to the request
req += '%s\r\n' % self.headers req += '%s\r\n' % self.headers
@ -270,7 +281,7 @@ class MitmProxy(HTTPServer):
def register_interceptor(self, interceptor_class): def register_interceptor(self, interceptor_class):
if not issubclass(interceptor_class, InterceptorPlugin): if not issubclass(interceptor_class, InterceptorPlugin):
raise InvalidInterceptorException('Expected type InterceptorPlugin got %s instead' % type(interceptor_class)) raise InvalidInterceptorPluginException('Expected type InterceptorPlugin got %s instead' % type(interceptor_class))
if issubclass(interceptor_class, RequestInterceptorPlugin): if issubclass(interceptor_class, RequestInterceptorPlugin):
self._req_plugins.append(interceptor_class) self._req_plugins.append(interceptor_class)
if issubclass(interceptor_class, ResponseInterceptorPlugin): if issubclass(interceptor_class, ResponseInterceptorPlugin):