ruleEngine.py

import re
import sys
import os
import logging
import time
import hashlib
from urlparse import urlparse
from malware import utils
from malware import decoder
from malware.snort import SnortRule
from malware.sql_tool import SQLiteTool
from malware import apikey
from virus_total_apis import PrivateApi as VirusTotal

logger = logging.getLogger(__name__)

REQUEST_RATE = 300
APIKEY = apikey.APIKEY_0


def clean_spaces(s):
    s = s.replace('\r', '')
    return s


class RuleEngineBase(object):
    def __init__(self, path='../PCAPLog/'):
        self.rules = list()
        self._db = SQLiteTool()
        self._db.creat_url_report()
        self.tcp_paylpad_iter = PayloadIterator2(path, 'tcp')
        self.udp_paylpad_iter = PayloadIterator2(path, 'udp')
        self.vd = Validator()
        self.vt = VirusTotal(APIKEY)

    def _make_rule(self, **kwargs):
        rule = SnortRule()
        rule.msg = '"Trojan.Gen"'

        content = kwargs.get('content')
        uricontent = kwargs.get('uricontent')
        dst_port = kwargs.get('dst_port')
        ref = kwargs.get('ref')
        protocol = kwargs.get('protocol')
        dst_port = kwargs.get('dst_port')

        if protocol is not None:
            rule.protocol = protocol
        if dst_port is not None:
            rule.dst_port = dst_port
        if content is not None:
            rule.content = content
        if uricontent is not None and uricontent != '/':
            rule.uricontent = uricontent
        if ref is not None:
            rule.ref = ref
        # pattern['sid'] = sid

        self.rules.append(rule)
        self._log_rules(rule, ref[0].split(',')[-1])

    def _get_url_positive(self, resource):
        urlkey = hashlib.sha1(resource).hexdigest()

        if self._db.is_key(urlkey):
            # print "In Table!!"
            return self._db.show_positive(urlkey)

    def _log_rules(self, data, filename):
        # print str(data)
        if not os.path.exists('./rules'):
            os.makedirs('./rules')

        with open('./rules/{m}_rule.rules'.format(m=filename), 'a') as fp:
            fp.write('{r}\n'.format(r=str(data)))


class RuleEngineOnline(RuleEngineBase):
    def __init__(self, path='../PCAPLog/'):
        self.vt_req_counter = 0
        self.vt_req_timer = time.time()
        super(RuleEngineOnline, self).__init__(path)

    def _check_timer_counter(self):
        if self.vt_req_counter == REQUEST_RATE:
            self.vt_req_counter = 0
            period = time.time() - self.vt_req_timer
            waiting = 60 - period + 1
            if waiting > 0:
                logger.info("Waiting %s seconds", (str(waiting)))
                time.sleep(waiting)
            self.vt_req_timer = time.time()

    def _make_rule(self, **kwargs):
        super(RuleEngineOnline, self)._make_rule(**kwargs)

    def _get_url_positive(self, resource):
        urlkey = hashlib.sha1(resource).hexdigest()

        if self._db.is_key(urlkey):
            # print "In Table!!"
            update_database = False
            if update_database:
                # ============== Updated the Database URL column ===============
                self._check_timer_counter()
                self.vt_req_counter += 1
                response = self.vt.get_url_report(resource)
                if response.get('error') is not None:
                    logger.info("Error: {e}".format(e=response.get('error')))
                    return None
                    # sys.exit(0)

                results = response.get('results')
                positives = results.get('positives')
                url = results.get('url')

                if positives >= 0:
                    self._db.insert2(urlkey, url, positives)
                # ============== Updated the Database URL column ===============
            return self._db.show_positive(urlkey)
        else:
            self._check_timer_counter()
            self.vt_req_counter += 1
            logger.info("Search on VirusTotal counter: %s",
                        str(self.vt_req_counter))
            logger.info(resource)
            response = self.vt.get_url_report(resource)

            if response.get('error') is not None:
                logger.info("Error: {e}".format(e=response.get('error')))
                return None
                # sys.exit(0)

            results = response.get('results')
            positives = results.get('positives')
            url = results.get('url')

            if positives >= 0:
                self._db.insert2(urlkey, url, positives)
                # self._db.insert2(url_id, url, positives)
                return positives
            elif positives is None:
                self._check_timer_counter()
                self.vt_req_counter += 1
                logger.info('''No report. Submmit the URL to VirusTotal countert: %s''',
                            str(self.vt_req_counter))
                self.vt.scan_url(resource)
                return None
            else:
                logger.debug("Get reports failed.")
                return None

    def _get_domain_positive(self, resource):
        domainkey = hashlib.sha1(resource).hexdigest()
        if self._db.is_key(domainkey):
            pass
            # return self._db.show_positive(urlkey)
        else:
            pass

    def http_rule_generate(self):
        for content, conn, filename in self.tcp_paylpad_iter:
            try:
                get_obj = self.vd.is_get_method(content)
                host_obj = self.vd.is_hsot(content)
                if host_obj and get_obj:
                    uri = get_obj.group(1)
                    host_field = clean_spaces(host_obj.group(1))
                    o = urlparse('http://'+ host_field + uri)
                    # domian = o.netloc
                    # uri = o.path
                    if o.path == '/':
                        # Proberbly an malicious domain name
                        domain_obj = self.vd.is_valid_url(host_field)
                        if domain_obj is not None:
                            domain_pos = self._get_url_positive(domain_obj.group(0))
                            if domain_pos > 0:
                                self._make_rule(protocol='tcp',
                                                content=['"{h}"'.format(h=clean_spaces(host_obj.group(0))), 'nocase'],
                                                dst_port=conn[3],
                                                ref=['md5,{m}'.format(m=filename.split('.')[0])])
                                                # md5=filename.split('.')[0])
                        else:
                            # Is a invalid url
                            pass

                    else:
                        # o.path != '/'
                        # string = self.vd.is_valid_utf8(host_field + uri)
                        # if string is not None:
                        #     Do search on VT

                        url_obj = self.vd.is_valid_url(host_field + uri)
                        if url_obj is not None:
                            url_pos = self._get_url_positive(url_obj.group(0))
                            if url_pos > 0:
                                self._make_rule(protocol='tcp',
                                                content=['"{h}"'.format(h=clean_spaces(host_obj.group(0))), 'nocase'],
                                                uricontent=['"{u}"'.format(u=o.path), 'nocase'],
                                                dst_port=conn[3],
                                                ref=['md5,{m}'.format(m=filename.split('.')[0])])
                                                # md5=filename.split('.')[0])
                        else:
                            # Is a invalid url
                            pass
                else:
                    pass
            except KeyboardInterrupt:
                logger.info("Quit")
                sys.exit()

    def dns_rule_generate(self):
        for content, conn, filename in self.udp_paylpad_iter:
            try:
                # print content, filename, conn[3]
                if content[0] == 'UNKNOWN_DNS':
                    # Bad DNS query opcode != 0
                    # print "Bad DNS query opcode != 0, %r" % content[1]
                    self._make_rule(protocol='udp',
                                    dst_port=conn[3],
                                    content=['"|'+content[1]+'|"'],
                                    ref=['md5,{m}'.format(m=filename.split('.')[0])])
                else:
                    domain_obj = self.vd.is_valid_url(content[0])
                    if domain_obj is not None:
                        domain_pos = self._get_url_positive(content[0])
                        if domain_pos > 0:
                            self._make_rule(protocol='udp',
                                            dst_port=conn[3],
                                            content=['"|'+content[1]+'|"'],
                                            ref=['md5,{m}'.format(m=filename.split('.')[0])])
                    else:
                        # Is a invalid domain name
                        with open('invalid_domain_name.log', 'a') as fp:
                            fp.write(filename+'\n')
                            fp.write(content[0]+'\n')
            except KeyboardInterrupt:
                logger.info("Quit")
                sys.exit()

    def _log_rules(self, data, filename):
        super(RuleEngineOnline, self)._log_rules(data, filename)


class Validator(object):
    def __init__(self):
        pass

    def is_valid_url(self, url):
        regex = re.compile(
            # r'^(?:[a-z0-9\.\-]*)://'  # scheme is validated separately
            r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}(?<!-)\.?)|'  # domain...
            r'localhost|'  # localhost...
            r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}|'  # ...or ipv4
            r'\[?[A-F0-9]*:[A-F0-9:]+\]?)'  # ...or ipv6
            r'(?::\d+)?'  # optional port
            r'(?:/?|[/?]\S+)$', re.IGNORECASE)
        return url is not None and regex.search(url)

    def is_valid_domain_name(self, domain_name):
        # TODO
        # Valid domain names
        # ex: syshell.exe is not domain
        # regex = re.compile(r'[a-zA-Z\d-]{,63}(\.[a-zA-Z\d-]{,63})*',
        #                    re.IGNORECASE)
        # return domain_name is not None and regex.search(domain_name)
        # return domain_name
        pass

    def is_hsot(self, content):
        regex = re.compile('Host: (.*)')
        return content is not None and regex.search(content)

    def is_get_method(self, content):
        regex = re.compile('GET (.*) ')
        return content is not None and regex.search(content)

    def is_valid_utf8(self, data):
        # valid_utf8 = True
        try:
            data.decode('utf-8')
            # return data
        except UnicodeDecodeError:
            with open('invalid_utf8.log', 'a') as fp:
                fp.write('{u}\n'.format(u=data))
            data = None
        return data
            # valid_utf8 = False


class PayloadIterator2(object):
    def __init__(self, path, protocol):
        self.index = 0
        self.path = path
        self.protocol = protocol
        self.content = list()
        self.five_tuple = list()
        self.file_pointer = list()

    def __iter__(self):
        pcap_list = list()
        for dirPath, dirNames, fileNames in os.walk(self.path):
            for f in fileNames:
                if f.endswith('.pcap'):
                    pcap_list.append(os.path.join(dirPath, f))
                else:
                    # Not a pcap file
                    pass

        if self.protocol == 'tcp':
            for p in pcap_list:
                connection = utils.follow_tcp_stream(p)
                for five_tuple, frame in connection.iteritems():
                    for seq, content in frame.iteritems():
                        if content:
                            # Generate the content and 5-tuple
                            self.content.append(content)
                            self.five_tuple.append(five_tuple)
                            self.file_pointer.append(p.split('/')[-1])
                        else:
                            # Some packets have no payload
                            pass
            logger.info("TCP Total Connections : %s",
                         str(len(set(self.five_tuple))))
        elif self.protocol == 'udp':
            for p in pcap_list:
                connection = decoder.decode_dns_qd_name(p)
                for five_tuple, qd_name_list in connection.iteritems():
                    self.content.append(qd_name_list)
                    self.five_tuple.append(five_tuple)
                    self.file_pointer.append(p.split('/')[-1])
            logger.info("UDP Total Connections : %s",
                         str(len(set(self.five_tuple))))
        else:
            logger.info("Protocol %s are not implement", self.protocol)

        logger.info("Total Pcap file: %s", str(len(set(pcap_list))))
        return self

    def next(self):
        try:
            five_tuple = self.five_tuple[self.index]
            content = self.content[self.index]
            file_pointer = self.file_pointer[self.index]
        except IndexError:
            raise StopIteration
        self.index += 1
        return content, five_tuple, file_pointer


def main():
    logging.basicConfig(level=logging.INFO,
                        format='[%(levelname)s] %(message)s',)
    udp_rules = list()
    tcp_rules = list()

    tcp_rule_engine = RuleEngineOnline()
    tcp_rule_engine.http_rule_generate()
    for ruleobj in tcp_rule_engine.rules:
        tcp_rules.append(str(ruleobj))

    udp_rule_engine = RuleEngineOnline()
    udp_rule_engine.dns_rule_generate()
    # print dir(rule_engine)
    for ruleobj in udp_rule_engine.rules:
        udp_rules.append(str(ruleobj))

    a = list(set(tcp_rules))
    for r in a:
        print r
        with open('tcp_snort.rules', 'a') as fp:
            fp.write(r + '\n')

    b = list(set(udp_rules))
    for r in b:
        print r
        with open('udp_snort.rules', 'a') as fp:
            fp.write(r + '\n')


if __name__ == "__main__":
    main()