LicenseMalwareBytes.py

# -*- coding: utf-8 -*-
import sys
import os
import stat
import errno
import socket
import ssl
import select
import http.client as httplib
import urllib.parse
import threading
import gzip
import zlib
import time
import json
import re
import logging
from http.server import HTTPServer, BaseHTTPRequestHandler
from socketserver import ThreadingMixIn
from io import StringIO
from html.parser import HTMLParser
import string
from OpenSSL import crypto
import argparse

from model import keystone

logger = None
CURRENT_DIR = os.path.dirname(os.path.abspath(sys.argv[0]))
CA_DIR = os.path.join(CURRENT_DIR, "CA")
if getattr(sys, 'frozen', False):
    CA_DIR = os.path.join(sys._MEIPASS, "CA")

class ThreadingHTTPServer(ThreadingMixIn, HTTPServer):
    address_family = socket.AF_INET
    daemon_threads = True

    def handle_error(self, request, client_address):
        # surpress socket/ssl related errors
        cls, e = sys.exc_info()[:2]
        if cls is socket.error or cls is ssl.SSLError:
            pass
        else:
            return HTTPServer.handle_error(self, request, client_address)


class ProxyRequestHandler(BaseHTTPRequestHandler):
    cakey = os.path.join(CA_DIR, "ca.key")
    cacert = os.path.join(CA_DIR, "ca.crt")
    certkey = os.path.join(CA_DIR, "client.key")
    certdir = os.path.join(os.environ["TMP"], "certs")
    timeout = 30
    lock = threading.Lock()

    def __init__(self, *args, **kwargs):
        self.tls = threading.local()
        self.tls.conns = {}
        self.ca_cert = crypto.load_certificate(crypto.FILETYPE_PEM, open(self.cacert, 'rb').read())

        self.ca_key = crypto.load_privatekey(crypto.FILETYPE_PEM, open(self.cakey, 'rb').read())

        self.cert_key = crypto.load_privatekey(crypto.FILETYPE_PEM, open(self.certkey, 'rb').read())
        BaseHTTPRequestHandler.__init__(self, *args, **kwargs)
        if not os.path.exists(self.certdir): os.makedirs(self.certdir)

    def log_error(self, format, *args):
        # surpress "Request timed out: timeout('timed out',)"
        if isinstance(args[0], socket.timeout):
            return

        self.log_message(format, *args)

    def do_CONNECT(self):
        if os.path.isfile(self.cakey) and os.path.isfile(self.cacert) and os.path.isfile(self.certkey) and os.path.isdir(self.certdir):
            self.connect_intercept()
        else:
            self.connect_relay()

    def connect_intercept(self):
        hostname = self.path.split(':')[0]
        certpath = os.path.join(self.certdir, hostname + ".crt")
        with self.lock:
            if not os.path.isfile(certpath):
                self.saveCertHostName(hostname, certpath)
        self.send_response(200, "Connection Established")
        self.end_headers()

        self.connection = ssl.wrap_socket(self.connection, keyfile=self.certkey, certfile=certpath, ca_certs=self.cacert, server_side=True)
        self.rfile = self.connection.makefile("rb", self.rbufsize)
        self.wfile = self.connection.makefile("wb", self.wbufsize)

        conntype = self.headers.get('Proxy-Connection', '')
        
        if self.protocol_version == "HTTP/1.1" and conntype.lower() != 'close':
            self.close_connection = 0
        else:
            self.close_connection = 1

    def connect_relay(self):
        address = self.path.split(':', 1)
        address[1] = int(address[1]) or 443
        try:
            s = socket.create_connection(address, timeout=self.timeout)
        except Exception as e:
            self.send_error(502)
            return
        self.send_response(200, 'Connection Established')
        self.end_headers()

        conns = [self.connection, s]
        self.close_connection = 0
        while not self.close_connection:
            rlist, wlist, xlist = select.select(conns, [], conns, self.timeout)
            if xlist or not rlist:
                break
            for r in rlist:
                other = conns[1] if r is conns[0] else conns[0]
                data = r.recv(8192)
                if not data:
                    self.close_connection = 1
                    break
                other.sendall(data)

    def saveCertHostName(self, hostname, certpath):
        epoch = int("%d" % (time.time() * 1000))
        certHost = crypto.X509()
        certHost.set_serial_number(epoch)
        certHost.set_issuer(self.ca_cert.get_subject())

        certHost_subj = certHost.get_subject()
        certHost_subj.commonName = hostname
        certHost.set_subject(certHost_subj)
        certHost.set_pubkey(self.cert_key)
        certHost.gmtime_adj_notBefore(0)
        certHost.gmtime_adj_notAfter(31536000)
        certHost.sign(self.ca_key, 'sha256')

        with open(certpath, "wb") as f:
            f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, certHost))

    def do_GET(self):

        if self.path == 'http://proxy2.test/':
            self.send_cacert()
            return
        req = self
        content_length = int(req.headers.get('Content-Length', 0))
        req_body = self.rfile.read(content_length) if content_length else None
       
        if req.path[0] == '/':
            if isinstance(self.connection, ssl.SSLSocket):
                req.path = "https://%s%s" % (req.headers['Host'], req.path)
            else:
                req.path = "http://%s%s" % (req.headers['Host'], req.path)

        if "keystone.mwbsys.com" in req.headers["Host"]:
            if req.command == "POST":
                if "/api/v1/installations/redeem.json" in req.path and req.headers.get("Content-Type", "") == "application/json":
                    body_json = json.loads(req_body)
                    keystone.installation_token = body_json["installation_token"]
                    keystone.license_key = body_json["license_key"]
                    data = keystone.getLicenseConfig()
                    self.send_resp(201, {"Content-Type": "application/json", "Content-Length": len(data)}, data)
                    return

                if "/api/v1/installations/check.json" in req.path and req.headers.get("Content-Type", "") == "application/json":
                    body_json = json.loads(req_body)
                    keystone.product_build = body_json["product_build"]
                    keystone.product_version = body_json["product_version"]
                    data = keystone.getLicenseConfig()
                    self.send_resp(201, {"Content-Type": "application/json", "Content-Length": len(data)}, data)
                    return

                if "/api/v1/installations/register.json" in req.path and req.headers.get("Content-Type", "") == "application/json":
                    body_json = json.loads(req_body)
                    keystone.product_build = body_json["product_build"]
                    keystone.product_version = body_json["product_version"]
                    keystone.machine_id = body_json["machine_id"]
                    data = keystone.getLicenseConfig()
                    self.send_resp(201, {"Content-Type": "application/json", "Content-Length": len(data)}, data)
                    return

        if "https://my-device.malwarebytes.com//devices" in req.path and req.command == "POST" and req.headers.get("Content-Type", "") == "application/json":
            body_json = json.loads(req_body)
            keystone.machine_id = body_json["machine_id"]
            keystone.name = body_json["name"]
            keystone.sku = body_json["sku"]
            data = '{"created":"ok"}'
            self.send_resp(201, {"Content-Type": "application/json", "Content-Length": len(data)}, data)
            return

        if "telemetry.malwarebytes.com" in req.headers["Host"]:
            if req.command == "POST":
                if "/api/v2/streams/applog/record" in req.path or "/api/v2/streams/user_actions/record" in req.path:
                    data = '{"status":"record added"}'
                    self.send_resp(201, {"Content-Type": "application/json", "Content-Length": len(data)}, data)
                    return

                if "/api/v2/streams/client/record" in req.path:
                    if  req.headers.get("Content-Type", "") == "application/json":
                        body_json = json.loads(req_body)
                        keystone.machine_id = body_json["header"]["machine_id"]
                    data = '{"status":"record added"}'
                    self.send_resp(201, {"Content-Type": "application/json", "Content-Length": len(data)}, data)
                    return

        req_body_modified = self.request_handler(req, req_body)
        if req_body_modified is False:
            self.send_error(403)
            return
        elif req_body_modified is not None:
            req_body = req_body_modified
            req.headers['Content-length'] = str(len(req_body))

        u = urllib.parse.urlsplit(req.path)
        scheme, netloc, path = u.scheme, u.netloc, (u.path + '?' + u.query if u.query else u.path)
        assert scheme in ('http', 'https')
        if netloc:
            del req.headers["Host"]
            req.headers['Host'] = netloc

        setattr(req, 'headers', self.filter_headers(req.headers))

        try:
            origin = (scheme, netloc)
            if not origin in self.tls.conns:
                if scheme == 'https':
                    self.tls.conns[origin] = httplib.HTTPSConnection(netloc, timeout=self.timeout, context=ssl._create_unverified_context())
                else:
                    self.tls.conns[origin] = httplib.HTTPConnection(netloc, timeout=self.timeout)
            conn = self.tls.conns[origin]
            conn.request(self.command, path, req_body, dict(req.headers))
            res = conn.getresponse()
            version_table = {10: 'HTTP/1.0', 11: 'HTTP/1.1'}
            setattr(res, 'headers', res.msg)
            setattr(res, 'response_version', version_table[res.version])

            # support streaming
            if not 'Content-Length' in res.headers and 'no-store' in res.headers.get('Cache-Control', ''):
                self.response_handler(req, req_body, res, '')
                setattr(res, 'headers', self.filter_headers(res.headers))
                self.relay_streaming(res)
                with self.lock:
                    self.save_handler(req, req_body, res, '')
                return

            res_body = res.read()
            #hexdump(res_body)
        except Exception as e:
            logger.debug(e)
            if origin in self.tls.conns:
                del self.tls.conns[origin]
            self.send_error(502)
            return

        content_encoding = res.headers.get('Content-Encoding', 'identity')
        res_body_plain = self.decode_content_body(res_body, content_encoding)

        res_body_modified = self.response_handler(req, req_body, res, res_body_plain)
        if res_body_modified is False:
            self.send_error(403)
            return
        elif res_body_modified is not None:
            print ("res_body_modified")
            res_body_plain = res_body_modified
            res_body = self.encode_content_body(res_body_plain, content_encoding)
            res.headers['Content-Length'] = str(len(res_body))

        setattr(res, 'headers', self.filter_headers(res.headers))

        self.send_resp(res.status, res.headers, res_body, res.reason)

        with self.lock:
            self.save_handler(req, req_body, res, res_body_plain)

    def send_resp(self, code, headers = {}, data = b"", message = None):
        self.send_response(code, message)
        for k, v in headers.items():
            self.send_header(k, v)
        self.end_headers()
        if data:
            if not isinstance(data, bytes): data = data.encode("utf-8")
            self.wfile.write(data)
            self.wfile.flush()
        
    def relay_streaming(self, res):
        self.send_response(res.status, res.reason)
        for k, v in res.headers.items():
            self.send_header(k, v)
        self.end_headers()
        try:
            while True:
                chunk = res.read(8192)
                if not chunk:
                    break
                self.wfile.write(chunk)
            self.wfile.flush()
        except socket.error:
            # connection closed by client
            pass

    do_HEAD = do_GET
    do_POST = do_GET
    do_PUT = do_GET
    do_DELETE = do_GET
    do_OPTIONS = do_GET
    do_PATCH = do_GET

    def filter_headers(self, headers):
        # http://tools.ietf.org/html/rfc2616#section-13.5.1
        hop_by_hop = ('connection', 'keep-alive', 'proxy-authenticate', 'proxy-authorization', 'te', 'trailers', 'transfer-encoding', 'upgrade')
        for k in hop_by_hop:
            del headers[k]

        # accept only supported encodings
        if 'Accept-Encoding' in headers:
            ae = headers['Accept-Encoding']
            filtered_encodings = [x for x in re.split(r',\s*', ae) if x in ('identity', 'gzip', 'x-gzip', 'deflate')]
            headers['Accept-Encoding'] = ', '.join(filtered_encodings)

        return headers

    def encode_content_body(self, text, encoding):
        if encoding == 'identity':
            data = text.encode("utf-8")
        elif encoding in ('gzip', 'x-gzip'):
            io = StringIO()
            with gzip.GzipFile(fileobj=io, mode='wb') as f:
                f.write(text)
            data = io.getvalue()
        elif encoding == 'deflate':
            data = zlib.compress(text)
        else:
            raise Exception("Unknown Content-Encoding: %s" % encoding)
        return data

    def decode_content_body(self, data, encoding):
        if encoding == 'identity':
            text = data
        elif encoding in ('gzip', 'x-gzip'):
            io = StringIO(data)
            with gzip.GzipFile(fileobj=io) as f:
                text = f.read()
        elif encoding == 'deflate':
            try:
                text = zlib.decompress(data)
            except zlib.error:
                text = zlib.decompress(data, -zlib.MAX_WBITS)
        else:
            raise Exception("Unknown Content-Encoding: %s" % encoding)
        return text

    def send_cacert(self):
        with open(self.cacert, 'rb') as f:
            data = f.read()
        self.send_resp(200, {"Content-Type": "application/x-x509-ca-cert", "Content-Length": len(data), "Connection": "close"}, data)

    def print_info(self, req, req_body, res, res_body):
        result = "\n"
        def parse_qsl(s):
            return '\n'.join("%-20s %s" % (k, v) for k, v in urllib.parse.parse_qsl(s, keep_blank_values=True))

        req_header_text = "%s %s %s\n%s" % (req.command, req.path, req.request_version, req.headers)
        res_header_text = "%s %d %s\n%s" % (res.response_version, res.status, res.reason, res.headers)

        result += req_header_text
        # print (with_color(33, req_header_text))

        u = urllib.parse.urlsplit(req.path)
        if u.query:
            query_text = parse_qsl(u.query)
            result += "==== QUERY PARAMETERS ====\n%s\n" % query_text
            # print (with_color(32, "==== QUERY PARAMETERS ====\n%s\n" % query_text))

        auth = req.headers.get('Authorization', '')
        if auth.lower().startswith('basic'):
            token = auth.split()[1].decode('base64')
            result += "==== BASIC AUTH ====\n%s\n" % token
            # print (with_color(31, "==== BASIC AUTH ====\n%s\n" % token))

        if req_body is not None:
            req_body_text = None
            content_type = req.headers.get('Content-Type', '')

            if content_type.startswith('application/x-www-form-urlencoded'):
                req_body_text = parse_qsl(req_body)
            elif content_type.startswith('application/json'):
                try:
                    json_obj = json.loads(req_body)
                    json_str = json.dumps(json_obj, indent=2)
                    if json_str.count('\n') < 50:
                        req_body_text = json_str
                    else:
                        lines = json_str.splitlines()
                        req_body_text = "%s\n(%d lines)" % ('\n'.join(lines[:50]), len(lines))
                except ValueError:
                    req_body_text = req_body
            elif len(req_body) < 1024:
                req_body_text = req_body

            if req_body_text:
                result += "==== REQUEST BODY ====\n%s\n" % req_body_text
                # print (with_color(32, "==== REQUEST BODY ====\n%s\n" % req_body_text))
        result += "\n"
        result += res_header_text
        # print (with_color(36, res_header_text))

        if res_body is not None:
            res_body_text = None
            content_type = res.headers.get('Content-Type', '')

            if content_type.startswith('application/json'):
                try:
                    json_obj = json.loads(res_body)
                    json_str = json.dumps(json_obj, indent=2)
                    if json_str.count('\n') < 50:
                        res_body_text = json_str
                    else:
                        lines = json_str.splitlines()
                        res_body_text = "%s\n(%d lines)" % ('\n'.join(lines[:50]), len(lines))
                except ValueError:
                    res_body_text = res_body
            elif content_type.startswith('text/html'):
                m = re.search(r'<title[^>]*>\s*([^<]+?)\s*</title>', res_body, re.I)
                if m:
                    h = HTMLParser()
                    result += "==== HTML TITLE ====\n%s\n" % h.unescape(m.group(1).decode('utf-8'))
                    # print (with_color(32, "==== HTML TITLE ====\n%s\n" % h.unescape(m.group(1).decode('utf-8'))))
            elif content_type.startswith('text/') and len(res_body) < 1024:
                res_body_text = res_body

            if res_body_text:
                result += "==== RESPONSE BODY ====\n%s\n" % res_body_text
                # print (with_color(32, "==== RESPONSE BODY ====\n%s\n" % res_body_text))
        logger.debug(result)

    def log_error(self, format, *args):
        logger.error(format % args)

    def log_message(self, format, *args):
        logger.debug("%s %s" % (self.address_string(), format % args))

    def request_handler(self, req, req_body):
        pass

    def response_handler(self, req, req_body, res, res_body):
        pass

    def save_handler(self, req, req_body, res, res_body):
        self.print_info(req, req_body, res, res_body)

def runProxy(serverAddress):
    ProxyRequestHandler.protocol_version = "HTTP/1.1"
    httpd = ThreadingHTTPServer(serverAddress, ProxyRequestHandler)
    sa = httpd.socket.getsockname()
    print (" +> Serving HTTP Proxy on", sa[0], "port", sa[1], "...")
    httpd.serve_forever()
    
def blockHost():
    hosts = ["keystone.mwbsys.com", "telemetry.malwarebytes.com"]
    hosts_file = os.path.join(os.environ["SystemRoot"], "System32", "drivers", "etc", "hosts")
    
    if os.path.exists(hosts_file):
        if not os.access(hosts_file, os.R_OK | os.W_OK):
            os.chmod(hosts_file, stat.S_IWRITE)
        try:
            with open(hosts_file, "r+") as f:
                f.seek(0, 0)
                data = f.read()
                hosts_tmp = hosts.copy()
                for host in hosts:
                    if host in data:
                        hosts_tmp.remove(host)
                if not hosts_tmp: return True
                f.seek(0, os.SEEK_END)
                f.write("\n")
                for host in hosts_tmp:
                    f.write("127.0.0.1 %s\n" % host)
                f.write("\n")
        except IOError as e:
            if e.errno == errno.EACCES:
                print (" +>",e)
                print (" +> Please remove read-only attribute in file", hosts_file)
                return False
            raise e
    else:
        with open(hosts_file, "w") as f:
            for host in hosts:
                f.write("127.0.0.1 %s\n" % host)
            f.write("\n")
    os.system('gpupdate /force')
    return True

def IsAdmin():
    try:
        return os.getuid() == 0
    except AttributeError:
        import ctypes
        return ctypes.windll.shell32.IsUserAnAdmin() != 0

def SetupLogger(loglevel = "info"):
    global logger
    logger = logging.getLogger(__name__)
    stdout_logger = logging.StreamHandler()
    file_logger = logging.FileHandler(os.path.basename(os.path.abspath(sys.argv[0])) + ".log", mode = 'w', encoding = "utf-8")
    formatter = logging.Formatter('%(asctime)s %(funcName)s %(levelname)s: %(message)s')
    stdout_logger.setFormatter(formatter)
    file_logger.setFormatter(formatter)
    logger.addHandler(stdout_logger)
    logger.addHandler(file_logger)
    logger.setLevel(logging._nameToLevel[loglevel.upper()])

def main():
   
    parser = argparse.ArgumentParser()
    parser.add_argument('--bind', '-b', default='', metavar='ADDRESS',
                        help='Specify alternate bind address '
                             '[default: all interfaces]')
    parser.add_argument('port', action='store',
                        default=9999, type=int,
                        nargs='?',
                        help='Specify alternate port [default: 9999]')
    parser.add_argument("-l", "--loglevel", metavar="LEVEL",
                        choices=["none", "error", "warning", "info", "debug"],
                        default="error",
                        help="""
                        Set the log message threshold.

                        Valid levels are: none, error, warning, info, debug
                        """
                        )
    args = parser.parse_args()
    if not IsAdmin():
        print (" +> Please run this program as an administrator.")
        sys.exit(0)
    if blockHost():
        print (" +> Block Host successed.")
    else:
        sys.exit(1)
    SetupLogger(args.loglevel)
    runProxy((args.bind, args.port))

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print ("Exit")
        sys.exit(0)