<commit_before>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)
        self.cache = {}

    def _make_rule(self, **kwargs):
        rule = SnortRule()
        pattern = dict()
        pattern['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:
            pattern['protocol'] = protocol
        if dst_port is not None:
            pattern['dst_port'] = dst_port
        if content is not None:
            pattern['content'] = content
        if uricontent is not None and uricontent != '/':
            pattern['uricontent'] = uricontent
        if ref is not None:
            pattern['reference'] = ref
        # pattern['sid'] = sid
        rule.set_malicious_pattern(**pattern)
        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.cache[urlkey] = [resource, positives]
                    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.cache[urlkey] = [resource, positives]
                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 None:
                            return

                        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:
                            # positive == 0
                            pass

                    else:
                        # o.path != '/'
                        # string = self.vd.is_valid_utf8(host_field + uri)
                        # if string is None:
                        #     return

                        url_obj = self.vd.is_valid_url(host_field + uri)
                        if url_obj is None:
                            return

                        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:
                            # positive == 0
                            pass
                else:
                    pass
            except KeyboardInterrupt:
                logger.info("Quit")
                sys.exit()
            # try:
            #     # print utils.connection_key_2_str(conn), filename
            #     # print content, utils.connection_key_2_str(conn)
            #     get_method = self.vd.is_get_method(content)
            #     host = self.vd.is_hsot(content)
            #     if host and get_method:
            #         if get_method.group(1) == '/':
            #             # Domain name
            #             url = self.vd.is_valid_url(host.group(1).rstrip())
            #         else:
            #             url = self.vd.is_valid_url(host.group(1).rstrip() +
            #                                        get_method.group(1))
            #
            #         # print url
            #         if url is not None:
            #             url = self.vd.is_valid_utf8(url.group(0))
            #
            #         if url is not None:
            #             host_content = host.group(0).rstrip()
            #             uricontent = get_method.group(1)
            #             pos = self._get_url_positive(url)
            #             if pos > 0:
            #                 if uricontent == '/':
            #                     uricontent = None
            #                 # print host_content
            #                 self._make_rule(content=host_content,
            #                                 uricontent=uricontent,
            #                                 dst_port=conn[3],
            #                                 md5=filename.split('.')[0])
            #                 # yield rule
            #             else:
            #                 # positives == 0 or positives == None
            #                 pass
            #         else:
            #             # 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 None:
                        with open('invalid_domain_name.log', 'a') as fp:
                            fp.write(filename+'\n')
                            fp.write(content[0]+'\n')
                        return
                    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:
                        pass
            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()
<commit_msg>[Bug] if got invalid url will return<commit_after>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()
        pattern = dict()
        pattern['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:
            pattern['protocol'] = protocol
        if dst_port is not None:
            pattern['dst_port'] = dst_port
        if content is not None:
            pattern['content'] = content
        if uricontent is not None and uricontent != '/':
            pattern['uricontent'] = uricontent
        if ref is not None:
            pattern['reference'] = ref
        # pattern['sid'] = sid
        rule.set_malicious_pattern(**pattern)
        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()
            # try:
            #     # print utils.connection_key_2_str(conn), filename
            #     # print content, utils.connection_key_2_str(conn)
            #     get_method = self.vd.is_get_method(content)
            #     host = self.vd.is_hsot(content)
            #     if host and get_method:
            #         if get_method.group(1) == '/':
            #             # Domain name
            #             url = self.vd.is_valid_url(host.group(1).rstrip())
            #         else:
            #             url = self.vd.is_valid_url(host.group(1).rstrip() +
            #                                        get_method.group(1))
            #
            #         # print url
            #         if url is not None:
            #             url = self.vd.is_valid_utf8(url.group(0))
            #
            #         if url is not None:
            #             host_content = host.group(0).rstrip()
            #             uricontent = get_method.group(1)
            #             pos = self._get_url_positive(url)
            #             if pos > 0:
            #                 if uricontent == '/':
            #                     uricontent = None
            #                 # print host_content
            #                 self._make_rule(content=host_content,
            #                                 uricontent=uricontent,
            #                                 dst_port=conn[3],
            #                                 md5=filename.split('.')[0])
            #                 # yield rule
            #             else:
            #                 # positives == 0 or positives == None
            #                 pass
            #         else:
            #             # 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()
<|endoftext|>


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

logger = logging.getLogger(__name__)

REQUEST_RATE = 300
APIKEY = apikey.APIKEY_0

class RuleEngineBase(object):
    def __init__(self, path='./PCAPLog/'):
        self.rules = list()
        self._db = sql_tool.SQLiteTool()
        self._db.creat()
        self.paylpad_iter = PayloadIterator2(path)
        self.vd = Validator()
        self.vt = VirusTotal(APIKEY)
        self.cache = {}

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

        content = kwargs.get('content')
        uricontent = kwargs.get('uricontent')
        dst_port = kwargs.get('dst_port')
        md5 = kwargs.get('md5')

        if content is not None:
            pattern['content'] = ['"{host}"'.format(host=content), 'nocase']
        if uricontent is not None and uricontent != '/':
            pattern['uricontent'] = ['"{uri}"'.format(uri=uricontent), 'nocase']
        if md5 is not None:
            pattern['reference'] = ['md5, {m}'.format(m=md5)]
        # pattern['sid'] = sid
        pattern['dst_port'] = dst_port
        rule.set_malicious_pattern(**pattern)
        self.rules.append(rule)
        self._log_rules(rule, md5)

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

        if self._db.is_url(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):
        self.vt_req_counter = 0
        self.vt_req_timer = time.time()
        super(RuleEngineOnline, self).__init__()

    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_url(urlkey):
            # print "In Table!!"
            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')
            # url_id = hashlib.sha1(url).hexdigest()

            if positives >= 0:
                # self.cache[urlkey] = [resource, positives]
                self._db.insert2(urlkey, resource, 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 generate(self):
        for content, conn, filename in self.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 = host_obj.group(1).rstrip()
                    o = urlparse('http://'+ host_field + uri)
                    # domian = o.netloc
                    # uri = o.path
                    if o.path != '/':
                        url = self.vd.is_valid_url(host_field + uri)
                        if url is None:
                            return
                        url = self.vd.is_valid_utf8(url)
                        if url is None:
                            return
                        url_pos = self._get_url_positive(url)
                        if url_pos > 0:
                            self._make_rule(content=host_obj.group(0).rstrip(),
                                            uricontent=o.path,
                                            dst_port=conn[3],
                                            md5=filename.split('.')[0])
                        else:
                            # positive == 0
                            pass
                    else:
                        pass
                        # o.path == '/'
                        # Proberbly an malicious domain name

                    # domain_pos = self._get_url_positive(o.netloc)
                    # # TODO
                    # # domain_pos = self._get_domain_positive(host_field)
                    # if domain_pos > 0:
                    #     self._make_rule(content=host_obj.group(0).rstrip(),
                    #                     dst_port=conn[3],
                    #                     md5=filename.split('.')[0])
                    # else:
                    #     # positive == 0
                    #     pass

                    # host_content = host_obj.group(0).rstrip()
                    # if uri == '/':
                    #     # Only Domain name
                    #     # domain == host_field
                    #     # d = self.vd.is_valid_utf8(host_field)
                    #     domain = host_field
                    #     domain_pos = self._get_url_positive(host_field)
                    #     # TODO
                    #     # domain_pos = self._get_domain_positive(host_field)
                    #     if domain_pos > 0:
                    #         # make content rules
                    #         self._make_rule(content=host_content,
                    #                         dst_port=conn[3],
                    #                         md5=filename.split('.')[0])
                    # else:
                    #     print('Domain + URI == URL')
                    #     # Domain + URI == URL
                    #     url = host_field + uri
                    #     # u = self.vd.is_valid_utf8(url)
                    #     url_obj = self.vd.is_valid_url(url)
                    #     if url_obj is None:
                    #         return
                    #     url_pos = self._get_url_positive(url)
                    #     if url and url_pos > 0:
                    #         # make content and uri rules
                    #         self._make_rule(content=host_content,
                    #                         uricontent=get_obj.group(1),
                    #                         dst_port=conn[3],
                    #                         md5=filename.split('.')[0])
                else:
                    pass
            except KeyboardInterrupt:
                logger.info("Quit")
                sys.exit()
            # try:
            #     # print utils.connection_key_2_str(conn), filename
            #     # print content, utils.connection_key_2_str(conn)
            #     get_method = self.vd.is_get_method(content)
            #     host = self.vd.is_hsot(content)
            #     if host and get_method:
            #         if get_method.group(1) == '/':
            #             # Domain name
            #             url = self.vd.is_valid_url(host.group(1).rstrip())
            #         else:
            #             url = self.vd.is_valid_url(host.group(1).rstrip() +
            #                                        get_method.group(1))
            #
            #         # print url
            #         if url is not None:
            #             url = self.vd.is_valid_utf8(url.group(0))
            #
            #         if url is not None:
            #             host_content = host.group(0).rstrip()
            #             uricontent = get_method.group(1)
            #             pos = self._get_url_positive(url)
            #             if pos > 0:
            #                 if uricontent == '/':
            #                     uricontent = None
            #                 # print host_content
            #                 self._make_rule(content=host_content,
            #                                 uricontent=uricontent,
            #                                 dst_port=conn[3],
            #                                 md5=filename.split('.')[0])
            #                 # yield rule
            #             else:
            #                 # positives == 0 or positives == None
            #                 pass
            #         else:
            #             # invalid_url
            #             pass
            #     else:
            #         pass
            # except KeyboardInterrupt:
            #     logger.info("Quit")
            #     sys.exit()

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

class RuleEngineOffline(RuleEngineBase):
    def __init__(self):
        super(RuleEngineOffline, self).__init__()

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

    def _log_rules(self, rule):
        super(RuleEngineOnline, self)._log_rules(rule)

    def generate(self):
        self.cache = pickle_tool.check_json()
        if not self.cache:
            print "Nothing can generate"
            sys.exit()
        for key, value in self.cache.iterms():
            if value[1] > 0:
                pass
                # TODO:
                # 1: JSON need to save uricontent, content
                # 2: Reconstruct the snort rule
                # get_method = self.vd.is_get_method(content)
                # host = self.vd.is_valid_url(content)
                # self._make_rule(content=, uricontent=, dst_port=80)

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):
        self.index = 0
        self.path = path
        self.pcap_list = list()
        self.content = list()
        self.five_tuple = list()
        self.file_pointer = list()

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

        for p in self.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("Total Pcap file: %s", str(len(set(self.file_pointer))))
        logger.info("Total Connections : %s", str(len(set(self.five_tuple))))
        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',)
    rules = list()

    rule_engine = RuleEngineOnline()
    rule_engine.generate()
    # print dir(rule_engine)

    for ruleobj in rule_engine.rules:
        rules.append(str(ruleobj))

    rules = list(set(rules))

    for r in rules:
        print r
        with open('main_snort.rules', 'a') as fp:
            fp.write(r + '\n')

if __name__ == "__main__":
    main()
<commit_msg>passed the string into self.vd.is_valid_utf8<commit_after>import re
import sys
import os
import logging
import time
import hashlib
from urlparse import urlparse
from malware import utils
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

class RuleEngineBase(object):
    def __init__(self, path='./PCAPLog/'):
        self.rules = list()
        self._db = SQLiteTool()
        self._db.creat()
        self.paylpad_iter = PayloadIterator2(path)
        self.vd = Validator()
        self.vt = VirusTotal(APIKEY)
        self.cache = {}

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

        content = kwargs.get('content')
        uricontent = kwargs.get('uricontent')
        dst_port = kwargs.get('dst_port')
        md5 = kwargs.get('md5')

        if content is not None:
            pattern['content'] = ['"{host}"'.format(host=content), 'nocase']
        if uricontent is not None and uricontent != '/':
            pattern['uricontent'] = ['"{uri}"'.format(uri=uricontent), 'nocase']
        if md5 is not None:
            pattern['reference'] = ['md5, {m}'.format(m=md5)]
        # pattern['sid'] = sid
        pattern['dst_port'] = dst_port
        rule.set_malicious_pattern(**pattern)
        self.rules.append(rule)
        self._log_rules(rule, md5)

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

        if self._db.is_url(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):
        self.vt_req_counter = 0
        self.vt_req_timer = time.time()
        super(RuleEngineOnline, self).__init__()

    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_url(urlkey):
            # print "In Table!!"
            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')
            # url_id = hashlib.sha1(url).hexdigest()

            if positives >= 0:
                # self.cache[urlkey] = [resource, positives]
                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 generate(self):
        for content, conn, filename in self.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 = host_obj.group(1).rstrip()
                    o = urlparse('http://'+ host_field + uri)
                    # domian = o.netloc
                    # uri = o.path
                    if o.path != '/':
                        string = self.vd.is_valid_utf8(host_field + uri)
                        if string is None:
                            return

                        url_obj = self.vd.is_valid_url(string)
                        if url_obj is None:
                            return

                        url_pos = self._get_url_positive(url_obj.group(0))
                        if url_pos > 0:
                            self._make_rule(content=host_obj.group(0).rstrip(),
                                            uricontent=o.path,
                                            dst_port=conn[3],
                                            md5=filename.split('.')[0])
                        else:
                            # positive == 0
                            pass
                    else:
                        pass
                        # o.path == '/'
                        # Proberbly an malicious domain name

                    # domain_pos = self._get_url_positive(o.netloc)
                    # # TODO
                    # # domain_pos = self._get_domain_positive(host_field)
                    # if domain_pos > 0:
                    #     self._make_rule(content=host_obj.group(0).rstrip(),
                    #                     dst_port=conn[3],
                    #                     md5=filename.split('.')[0])
                    # else:
                    #     # positive == 0
                    #     pass

                    # host_content = host_obj.group(0).rstrip()
                    # if uri == '/':
                    #     # Only Domain name
                    #     # domain == host_field
                    #     # d = self.vd.is_valid_utf8(host_field)
                    #     domain = host_field
                    #     domain_pos = self._get_url_positive(host_field)
                    #     # TODO
                    #     # domain_pos = self._get_domain_positive(host_field)
                    #     if domain_pos > 0:
                    #         # make content rules
                    #         self._make_rule(content=host_content,
                    #                         dst_port=conn[3],
                    #                         md5=filename.split('.')[0])
                    # else:
                    #     print('Domain + URI == URL')
                    #     # Domain + URI == URL
                    #     url = host_field + uri
                    #     # u = self.vd.is_valid_utf8(url)
                    #     url_obj = self.vd.is_valid_url(url)
                    #     if url_obj is None:
                    #         return
                    #     url_pos = self._get_url_positive(url)
                    #     if url and url_pos > 0:
                    #         # make content and uri rules
                    #         self._make_rule(content=host_content,
                    #                         uricontent=get_obj.group(1),
                    #                         dst_port=conn[3],
                    #                         md5=filename.split('.')[0])
                else:
                    pass
            except KeyboardInterrupt:
                logger.info("Quit")
                sys.exit()
            # try:
            #     # print utils.connection_key_2_str(conn), filename
            #     # print content, utils.connection_key_2_str(conn)
            #     get_method = self.vd.is_get_method(content)
            #     host = self.vd.is_hsot(content)
            #     if host and get_method:
            #         if get_method.group(1) == '/':
            #             # Domain name
            #             url = self.vd.is_valid_url(host.group(1).rstrip())
            #         else:
            #             url = self.vd.is_valid_url(host.group(1).rstrip() +
            #                                        get_method.group(1))
            #
            #         # print url
            #         if url is not None:
            #             url = self.vd.is_valid_utf8(url.group(0))
            #
            #         if url is not None:
            #             host_content = host.group(0).rstrip()
            #             uricontent = get_method.group(1)
            #             pos = self._get_url_positive(url)
            #             if pos > 0:
            #                 if uricontent == '/':
            #                     uricontent = None
            #                 # print host_content
            #                 self._make_rule(content=host_content,
            #                                 uricontent=uricontent,
            #                                 dst_port=conn[3],
            #                                 md5=filename.split('.')[0])
            #                 # yield rule
            #             else:
            #                 # positives == 0 or positives == None
            #                 pass
            #         else:
            #             # invalid_url
            #             pass
            #     else:
            #         pass
            # except KeyboardInterrupt:
            #     logger.info("Quit")
            #     sys.exit()

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

class RuleEngineOffline(RuleEngineBase):
    def __init__(self):
        super(RuleEngineOffline, self).__init__()

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

    def _log_rules(self, rule):
        super(RuleEngineOnline, self)._log_rules(rule)

    def generate(self):
        self.cache = pickle_tool.check_json()
        if not self.cache:
            print "Nothing can generate"
            sys.exit()
        for key, value in self.cache.iterms():
            if value[1] > 0:
                pass
                # TODO:
                # 1: JSON need to save uricontent, content
                # 2: Reconstruct the snort rule
                # get_method = self.vd.is_get_method(content)
                # host = self.vd.is_valid_url(content)
                # self._make_rule(content=, uricontent=, dst_port=80)

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):
        self.index = 0
        self.path = path
        self.pcap_list = list()
        self.content = list()
        self.five_tuple = list()
        self.file_pointer = list()

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

        for p in self.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("Total Pcap file: %s", str(len(set(self.file_pointer))))
        logger.info("Total Connections : %s", str(len(set(self.five_tuple))))
        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',)
    rules = list()

    rule_engine = RuleEngineOnline()
    rule_engine.generate()
    # print dir(rule_engine)

    for ruleobj in rule_engine.rules:
        rules.append(str(ruleobj))

    rules = list(set(rules))

    for r in rules:
        print r
        with open('main_snort.rules', 'a') as fp:
            fp.write(r + '\n')

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>"""
DESCRIPTION
-----------
    For each directory provided on the command line create a table
    in that directory with one row for each FITS file in the directory.
    The columns are FITS keywords extracted from the header of each
    file.

    The list of default keywords extracted is available through the command
    line option ``--list-default``.

    .. Note::
        This feature is available only from the command line.

    For more control over the parameters see :func:`triage_fits_files`

    .. Note::
        This script is **NOT RECURSIVE**; it will not process files in
        subdirectories of the the directories supplied on the command line.

EXAMPLES
--------

    Invoking this script from the command line::

        python run_triage.py /my/folder/of/images

    Get list of default keywords included in summary table::

        python run_triage.py --list-default

    To work on the same folder from within python, do this::

        from run_triage import triage_directories
        astrometry_for_directory('/my/folder/of/images')

"""
import os
import numpy as np

# The import below
# MUST happen before any logging...oddly, setting the VALUE
# doesn't change anything. True story.

from astropy.logger import LOG_WARNINGS
import logging
from customlogger import console_handler, add_file_handlers

logger = logging.getLogger()
screen_handler = console_handler()
logger.addHandler(screen_handler)

class DefaultFileNames(object):
    def __init__(self):
        self.object_file_name = 'NEEDS_OBJECT_NAME.txt'
        self.pointing_file_name = 'NEEDS_POINTING_INFO.txt'
        self.filter_file_name = 'NEEDS_FILTER.txt'
        self.output_table = 'Manifest.txt'

    def as_dict(self):
        return self.__dict__

def write_list(dir, file, info, column_name='File'):
    from astropy.table import Table
    temp_table = Table(data=[info],
                       names=[column_name])
    temp_table.write(os.path.join(dir, file),
                     format='ascii')

def contains_maximdl_imagetype(image_collection):
    """
    Check an image file collection for MaxImDL-style image types
    """
    import re
    file_info = image_collection.summary_info
    image_types = ' '.join([typ for typ in file_info['imagetyp']])
    if re.search('[fF]rame', image_types) is not None:
        return True
    else:
        return False

def triage_fits_files(dir='.', file_info_to_keep=['imagetyp',
                                                  'object',
                                                  'filter']):
    """
    Check FITS files in a directory for deficient headers

    `dir` is the name of the directory to search for files.

    `file_info_to_keep` is a list of the FITS keywords to get values
    for for each FITS file in `dir`.
    """
    from feder import Feder
    from image_collection import ImageFileCollection

    all_file_info = file_info_to_keep
    feder = Feder()
    RA = feder.RA
    if 'ra' not in [key.lower() for key in all_file_info]:
        all_file_info.extend(RA.names)

    images = ImageFileCollection(dir, keywords=all_file_info)
    file_info = images.fits_summary(keywords=all_file_info)

    # check for bad image type and halt until that is fixed.
    if contains_maximdl_imagetype(images):
        raise ValueError(
            'Correct MaxImDL-style image types before proceeding.')

    file_needs_filter = \
        list(images.files_filtered(imagetyp='light',
                                   filter=None))
    file_needs_filter += \
        list(images.files_filtered(imagetyp='flat',
                                   filter=None))

    file_needs_object_name = \
        list(images.files_filtered(imagetyp='light',
                                   object=None))

    lights = file_info[file_info['imagetyp'] == 'LIGHT']
    has_no_ra = np.array([True] * len(lights))
    for ra_name in RA.names:
        try:
            has_no_ra &= (lights[ra_name].mask)
        except KeyError as e:
            pass

    needs_minimal_pointing = (lights['object'].mask) & has_no_ra

    dir_info = {'files': file_info,
                'needs_filter': file_needs_filter,
                'needs_pointing': list(lights['file'][needs_minimal_pointing]),
                'needs_object_name': file_needs_object_name}
    return dir_info

def triage_directories(directories,
                       keywords=None,
                       object_file_name=None,
                       pointing_file_name=None,
                       filter_file_name=None,
                       output_table=None,
                       destination=None):

    use_keys = []
    if keywords is not None:
        use_keys = keywords

    for currentDir in directories:
        logger.info('Examining directory %s', currentDir)
        result = triage_fits_files(currentDir, file_info_to_keep=use_keys)
        outfiles = [pointing_file_name, filter_file_name,
                    object_file_name, output_table]
        for fil in [outfile for outfile in outfiles if outfile is not None]:
            try:
                os.remove(os.path.join(currentDir, fil))
            except OSError:
                pass

        need_pointing = result['needs_pointing']
        need_filter = result['needs_filter']
        need_object_name = result['needs_object_name']
        if destination is not None:
            target_dir = destination
        else:
            target_dir = currentDir
        if need_pointing and pointing_file_name is not None:
            write_list(target_dir, pointing_file_name, need_pointing)
        if need_filter and filter_file_name is not None:
            write_list(target_dir, filter_file_name, need_filter)
        if need_object_name and object_file_name is not None:
            write_list(target_dir, object_file_name,
                       need_object_name)
        tbl = result['files']
        if ((len(tbl) > 0) and (output_table is not None)):
            tbl.write(os.path.join(target_dir, output_table),
                      format='ascii', delimiter=',')

def construct_parser():
    from argparse import ArgumentParser
    import script_helpers

    parser = ArgumentParser()
    script_helpers.setup_parser_help(parser, __doc__)
    # allow for no directories below so that -l option
    # can be used without needing to specify a directory
    script_helpers.add_directories(parser, '*')
    script_helpers.add_debug(parser)
    script_helpers.add_verbose(parser)
    script_helpers.add_destination_directory(parser)
    script_helpers.add_no_log_destination(parser)

    key_help = 'FITS keyword to add to table in addition to the defaults; '
    key_help += 'for multiple keywords use this option multiple times.'
    parser.add_argument('-k', '--key', action='append',
                        help=key_help)

    no_default_help = 'Do not include default list of keywords in table'
    parser.add_argument('--no-default', action='store_true',
                        help=no_default_help)

    list_help = 'Print default list keywords put into table and exit'
    parser.add_argument('-l', '--list-default', action='store_true',
                        help=list_help)

    default_names = DefaultFileNames()
    output_file_help = 'Name of file in which table is saved; default is '
    output_file_help += default_names.output_table
    parser.add_argument('-t', '--table-name',
                        default=default_names.output_table,
                        help=output_file_help)

    needs_object_help = 'Name of file to which list of files that need '
    needs_object_help += 'object name is saved; default is '
    needs_object_help += default_names.object_file_name
    parser.add_argument('-o', '--object-needed-list',
                        default=default_names.object_file_name,
                        help=needs_object_help)

    needs_pointing_help = 'Name of file to which list of files that need '
    needs_pointing_help += 'pointing name is saved; default is '
    needs_pointing_help += default_names.pointing_file_name
    parser.add_argument('-p', '--pointing-needed-list',
                        default=default_names.pointing_file_name,
                        help=needs_pointing_help)

    needs_filter_help = 'Name of file to which list of files that need '
    needs_filter_help += 'filter is saved; default is '
    needs_filter_help += default_names.filter_file_name
    parser.add_argument('-f', '--filter-needed-list',
                        default=default_names.filter_file_name,
                        help=needs_filter_help)

    return parser

always_include_keys = ['imagetyp', 'filter', 'exptime', 'ccd-temp',
                       'object', 'observer', 'airmass', 'instrument',
                       'RA', 'Dec']

if __name__ == "__main__":
    from sys import exit
    from os import getcwd
    from script_helpers import setup_logging

    parser = construct_parser()
    args = parser.parse_args()

    #all_keywords = ['imagetyp', 'filter', 'exptime', 'ccd-temp']

    setup_logging(logger, args, screen_handler)

    add_file_handlers(logger, getcwd(), 'run_triage')

    try:
        always_include_keys.extend(args.key)
    except TypeError as e:
        pass

    if args.no_default:
        always_include_keys = None

    if args.list_default:
        print 'Keys included by default are:\n'
        keys_print = [key.upper() for key in always_include_keys]
        print ', '.join(keys_print)
        exit(0)

    if not args.dir:
        parser.error('No directory specified')

    triage_directories(args.dir, keywords=always_include_keys,
                       object_file_name=args.object_needed_list,
                       pointing_file_name=args.pointing_needed_list,
                       filter_file_name=args.filter_needed_list,
                       output_table=args.table_name,
                       destination=args.destination_dir)
<commit_msg>Fix logging in destination directory<commit_after>"""
DESCRIPTION
-----------
    For each directory provided on the command line create a table
    in that directory with one row for each FITS file in the directory.
    The columns are FITS keywords extracted from the header of each
    file.

    The list of default keywords extracted is available through the command
    line option ``--list-default``.

    .. Note::
        This feature is available only from the command line.

    For more control over the parameters see :func:`triage_fits_files`

    .. Note::
        This script is **NOT RECURSIVE**; it will not process files in
        subdirectories of the the directories supplied on the command line.

EXAMPLES
--------

    Invoking this script from the command line::

        python run_triage.py /my/folder/of/images

    Get list of default keywords included in summary table::

        python run_triage.py --list-default

    To work on the same folder from within python, do this::

        from run_triage import triage_directories
        astrometry_for_directory('/my/folder/of/images')

"""
import os
import numpy as np

# The import below
# MUST happen before any logging...oddly, setting the VALUE
# doesn't change anything. True story.

from astropy.logger import LOG_WARNINGS
import logging
from customlogger import console_handler, add_file_handlers

logger = logging.getLogger()
screen_handler = console_handler()
logger.addHandler(screen_handler)

class DefaultFileNames(object):
    def __init__(self):
        self.object_file_name = 'NEEDS_OBJECT_NAME.txt'
        self.pointing_file_name = 'NEEDS_POINTING_INFO.txt'
        self.filter_file_name = 'NEEDS_FILTER.txt'
        self.output_table = 'Manifest.txt'

    def as_dict(self):
        return self.__dict__

def write_list(dir, file, info, column_name='File'):
    from astropy.table import Table
    temp_table = Table(data=[info],
                       names=[column_name])
    temp_table.write(os.path.join(dir, file),
                     format='ascii')

def contains_maximdl_imagetype(image_collection):
    """
    Check an image file collection for MaxImDL-style image types
    """
    import re
    file_info = image_collection.summary_info
    image_types = ' '.join([typ for typ in file_info['imagetyp']])
    if re.search('[fF]rame', image_types) is not None:
        return True
    else:
        return False

def triage_fits_files(dir='.', file_info_to_keep=['imagetyp',
                                                  'object',
                                                  'filter']):
    """
    Check FITS files in a directory for deficient headers

    `dir` is the name of the directory to search for files.

    `file_info_to_keep` is a list of the FITS keywords to get values
    for for each FITS file in `dir`.
    """
    from feder import Feder
    from image_collection import ImageFileCollection

    all_file_info = file_info_to_keep
    feder = Feder()
    RA = feder.RA
    if 'ra' not in [key.lower() for key in all_file_info]:
        all_file_info.extend(RA.names)

    images = ImageFileCollection(dir, keywords=all_file_info)
    file_info = images.fits_summary(keywords=all_file_info)

    # check for bad image type and halt until that is fixed.
    if contains_maximdl_imagetype(images):
        raise ValueError(
            'Correct MaxImDL-style image types before proceeding.')

    file_needs_filter = \
        list(images.files_filtered(imagetyp='light',
                                   filter=None))
    file_needs_filter += \
        list(images.files_filtered(imagetyp='flat',
                                   filter=None))

    file_needs_object_name = \
        list(images.files_filtered(imagetyp='light',
                                   object=None))

    lights = file_info[file_info['imagetyp'] == 'LIGHT']
    has_no_ra = np.array([True] * len(lights))
    for ra_name in RA.names:
        try:
            has_no_ra &= (lights[ra_name].mask)
        except KeyError as e:
            pass

    needs_minimal_pointing = (lights['object'].mask) & has_no_ra

    dir_info = {'files': file_info,
                'needs_filter': file_needs_filter,
                'needs_pointing': list(lights['file'][needs_minimal_pointing]),
                'needs_object_name': file_needs_object_name}
    return dir_info

def triage_directories(directories,
                       keywords=None,
                       object_file_name=None,
                       pointing_file_name=None,
                       filter_file_name=None,
                       output_table=None,
                       destination=None,
                       no_log_destination=False):

    use_keys = []
    if keywords is not None:
        use_keys = keywords

    for currentDir in directories:
        if destination is not None:
            target_dir = destination
        else:
            target_dir = currentDir
        if not no_log_destination:
            add_file_handlers(logger, destination, 'run_triage')
        logger.info('Examining directory %s', currentDir)
        result = triage_fits_files(currentDir, file_info_to_keep=use_keys)
        outfiles = [pointing_file_name, filter_file_name,
                    object_file_name, output_table]
        for fil in [outfile for outfile in outfiles if outfile is not None]:
            try:
                os.remove(os.path.join(currentDir, fil))
            except OSError:
                pass

        need_pointing = result['needs_pointing']
        need_filter = result['needs_filter']
        need_object_name = result['needs_object_name']

        if need_pointing and pointing_file_name is not None:
            write_list(target_dir, pointing_file_name, need_pointing)
        if need_filter and filter_file_name is not None:
            write_list(target_dir, filter_file_name, need_filter)
        if need_object_name and object_file_name is not None:
            write_list(target_dir, object_file_name,
                       need_object_name)
        tbl = result['files']
        if ((len(tbl) > 0) and (output_table is not None)):
            tbl.write(os.path.join(target_dir, output_table),
                      format='ascii', delimiter=',')

def construct_parser():
    from argparse import ArgumentParser
    import script_helpers

    parser = ArgumentParser()
    script_helpers.setup_parser_help(parser, __doc__)
    # allow for no directories below so that -l option
    # can be used without needing to specify a directory
    script_helpers.add_directories(parser, '*')
    script_helpers.add_debug(parser)
    script_helpers.add_verbose(parser)
    script_helpers.add_destination_directory(parser)
    script_helpers.add_no_log_destination(parser)
    script_helpers.add_console_output_args(parser)

    key_help = 'FITS keyword to add to table in addition to the defaults; '
    key_help += 'for multiple keywords use this option multiple times.'
    parser.add_argument('-k', '--key', action='append',
                        help=key_help)

    no_default_help = 'Do not include default list of keywords in table'
    parser.add_argument('--no-default', action='store_true',
                        help=no_default_help)

    list_help = 'Print default list keywords put into table and exit'
    parser.add_argument('-l', '--list-default', action='store_true',
                        help=list_help)

    default_names = DefaultFileNames()
    output_file_help = 'Name of file in which table is saved; default is '
    output_file_help += default_names.output_table
    parser.add_argument('-t', '--table-name',
                        default=default_names.output_table,
                        help=output_file_help)

    needs_object_help = 'Name of file to which list of files that need '
    needs_object_help += 'object name is saved; default is '
    needs_object_help += default_names.object_file_name
    parser.add_argument('-o', '--object-needed-list',
                        default=default_names.object_file_name,
                        help=needs_object_help)

    needs_pointing_help = 'Name of file to which list of files that need '
    needs_pointing_help += 'pointing name is saved; default is '
    needs_pointing_help += default_names.pointing_file_name
    parser.add_argument('-p', '--pointing-needed-list',
                        default=default_names.pointing_file_name,
                        help=needs_pointing_help)

    needs_filter_help = 'Name of file to which list of files that need '
    needs_filter_help += 'filter is saved; default is '
    needs_filter_help += default_names.filter_file_name
    parser.add_argument('-f', '--filter-needed-list',
                        default=default_names.filter_file_name,
                        help=needs_filter_help)

    return parser

always_include_keys = ['imagetyp', 'filter', 'exptime', 'ccd-temp',
                       'object', 'observer', 'airmass', 'instrume',
                       'RA', 'Dec']

if __name__ == "__main__":
    from sys import exit
    from os import getcwd
    from script_helpers import setup_logging

    parser = construct_parser()
    args = parser.parse_args()

    #all_keywords = ['imagetyp', 'filter', 'exptime', 'ccd-temp']

    setup_logging(logger, args, screen_handler)

    add_file_handlers(logger, getcwd(), 'run_triage')

    try:
        always_include_keys.extend(args.key)
    except TypeError as e:
        pass

    if args.no_default:
        always_include_keys = None

    if args.list_default:
        print 'Keys included by default are:\n'
        keys_print = [key.upper() for key in always_include_keys]
        print ', '.join(keys_print)
        exit(0)

    if not args.dir:
        parser.error('No directory specified')

    triage_directories(args.dir, keywords=always_include_keys,
                       object_file_name=args.object_needed_list,
                       pointing_file_name=args.pointing_needed_list,
                       filter_file_name=args.filter_needed_list,
                       output_table=args.table_name,
                       destination=args.destination_dir)
<|endoftext|>


<commit_before>import sys

from difflib import SequenceMatcher

from pygments import lex
from pygments.lexers.c_cpp import CLexer
from pygments.styles import get_style_by_name

class CodeDiff:
  opcode_style = {
    'insert': {'bgcolor': 'eaffea'}, 'replace': {'bgcolor': 'fff68f'},
    'delete': {'bgcolor': 'ffaaaa'}, 'equal': {'bgcolor': None},
  }

  def __init__(self, from_str, to_str):
    self.from_str = from_str
    self.to_str = to_str

    # difflib stuff.
    self.sm = SequenceMatcher(None, from_str, to_str)

    # Pygments stuff.
    # TODO: allow chosing of style.
    self.style = get_style_by_name('colorful')

    self.style_map = {}
    for token_type, style in self.style:
      self.style_map[token_type] = style

  def get_desc(self, for_from_str):
    descs = []
    pdesc = self.get_pygmentation_desc(for_from_str)
    ddesc = self.get_difflib_desc(for_from_str)
    while pdesc or ddesc:
      # The diff may be larger than the program itself.
      if not pdesc:
        desc = ddesc.pop(0)

        # Update the starting position of the diff fragment to the last
        # character descriptor in descriptor list.
        desc[1] = descs[-1][2]
        descs.append(desc)
        break

      # We have an opcode from dstart to dend.
      dstyle, dstart, dend = ddesc[0]

      # We have an token dstart to dend.
      pstyle, pstart, pend = pdesc[0]

      # TODO: update() overwrites values for common keys, is this correct?
      style = pstyle.copy()
      style.update(dstyle)

      # If the token fits within the opcode region, merge the two styles.
      if dstart <= pstart and pend <= dend:
        descs.append([style, pstart, pend])

        # Next token & remove the diff opcode if necessary.
        pdesc.pop(0)
        if (pend == dend): ddesc.pop(0)

      # If the token doesn't fit within the opcode region, seperate it into two.
      elif dstart <= pstart and pend >= dend:
        descs.append([style, pstart, dend])

        # Second part of the token stays in the array.
        pdesc[0] = [style, dend, pend]
        ddesc.pop(0)
      else:
        # TODO: missing some cases?
        assert False

    # Both arrays should be empty when the above loop finishes.
    assert (not pdesc) and (not ddesc)
    return descs

  def get_pygmentation_desc(self, for_from_str):
    desc = []
    s = self.from_str if for_from_str else self.to_str
    token_start = token_end = 0
    for token_type, token in lex(s, CLexer()):
      token_start = token_end
      token_end += len(token)
      desc.append([self.style_map[token_type], token_start, token_end])
    return desc

  def get_difflib_desc(self, for_from_str):
    desc = []
    for s in self.sm.get_opcodes():
      opcode, from_start, from_end, to_start, to_end = s
      if for_from_str:
        desc.append([CodeDiff.opcode_style[opcode], from_start, from_end])
      else:
        desc.append([CodeDiff.opcode_style[opcode], to_start, to_end])
    return desc

  def to_html(self, for_from_str):
    # TODO: don't hardcode output formatting.
    html = ""
    descs = self.get_desc(for_from_str)
    for style, start, end in descs:
      span = ''
      fragment = (self.from_str if for_from_str else self.to_str)[start:end]
      fragment = fragment.replace('\n', '
')
      fragment = fragment.replace(' ', ' ')
      span += fragment + ''
      html += span
    return html

  @staticmethod
  def from_files(from_file, to_file):
    with open(from_file) as f:
      from_str = f.read()
    with open(to_file) as f:
      to_str = f.read()
    return CodeDiff(from_str, to_str)

<commit_msg>Style with classes in HTML output.<commit_after>import sys

from difflib import SequenceMatcher

from pygments import lex
from pygments.lexers.c_cpp import CLexer
from pygments.styles import get_style_by_name

class CodeDiff:
  opcode_style = {
    'insert': {'bgcolor': 'eaffea'}, 'replace': {'bgcolor': 'fff68f'},
    'delete': {'bgcolor': 'ffaaaa'}, 'equal': {'bgcolor': None},
  }

  def __init__(self, from_str, to_str):
    self.from_str = from_str
    self.to_str = to_str

    # difflib stuff.
    self.sm = SequenceMatcher(None, from_str, to_str)

  def get_desc(self, for_from_str):
    descs = []
    pdesc = self.get_pygmentation_desc(for_from_str)
    ddesc = self.get_difflib_desc(for_from_str)
    while pdesc or ddesc:
      # The diff may be larger than the program itself.
      if not pdesc:
        desc = ddesc.pop(0)

        # Update the starting position of the diff fragment to the last
        # character descriptor in descriptor list.
        desc[1] = descs[-1][2]
        descs.append(desc)
        break

      # We have an opcode from dstart to dend.
      dstyle, dstart, dend = ddesc[0]

      # We have an token dstart to dend.
      pstyle, pstart, pend = pdesc[0]

      # TODO: update() overwrites values for common keys, is this correct?
      style = pstyle + dstyle

      # If the token fits within the opcode region, merge the two styles.
      if dstart <= pstart and pend <= dend:
        descs.append([style, pstart, pend])

        # Next token & remove the diff opcode if necessary.
        pdesc.pop(0)
        if (pend == dend): ddesc.pop(0)

      # If the token doesn't fit within the opcode region, seperate it into two.
      elif dstart <= pstart and pend >= dend:
        descs.append([style, pstart, dend])

        # Second part of the token stays in the array.
        pdesc[0] = [style, dend, pend]
        ddesc.pop(0)
      else:
        # TODO: missing some cases?
        assert False

    # Both arrays should be empty when the above loop finishes.
    assert (not pdesc) and (not ddesc)
    return descs

  def get_pygmentation_desc(self, for_from_str):
    desc = []
    s = self.from_str if for_from_str else self.to_str
    token_start = token_end = 0
    for token_type, token in lex(s, CLexer()):
      token_start = token_end
      token_end += len(token)
      token_class = 'token-' + '-'.join(token_type).lower()
      desc.append([[token_class], token_start, token_end])
    return desc

  def get_difflib_desc(self, for_from_str):
    desc = []
    for s in self.sm.get_opcodes():
      opcode, from_start, from_end, to_start, to_end = s
      if for_from_str:
        desc.append([['opcode-' + opcode], from_start, from_end])
      else:
        desc.append([['opcode-' + opcode], to_start, to_end])
    return desc

  def to_html(self, for_from_str):
    # TODO: don't hardcode output formatting.
    html = ""
    descs = self.get_desc(for_from_str)
    for style, start, end in descs:
      span = ''
      fragment = (self.from_str if for_from_str else self.to_str)[start:end]
      fragment = fragment.replace('\n', '
')
      fragment = fragment.replace(' ', ' ')
      span += fragment + ''
      html += span
    return html

  @staticmethod
  def from_files(from_file, to_file):
    with open(from_file) as f:
      from_str = f.read()
    with open(to_file) as f:
      to_str = f.read()
    return CodeDiff(from_str, to_str)

<|endoftext|>


<commit_before>import uuid
from datetime import timedelta
from datetime import time

from django.conf import settings
from django.utils.decorators import method_decorator
from django.views.decorators.csrf import ensure_csrf_cookie
from django.views.generic.base import TemplateView
from django.utils.safestring import mark_safe

from .models import *

import redis
# TODO(arvnd): Should probably get the redis adapter from some common
# place.
REDIS = redis.StrictRedis(host=settings.MMVP_REDIS_HOST,
                          port=settings.MMVP_REDIS_PORT,
                          db=settings.MMVP_REDIS_DB,
                          charset="utf-8",
                          decode_responses=True)

def sizeof_fmt(num, suffix='B'):
    for unit in ['','K','M','G']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'T', suffix)

# Darjeeling Limited
progress_bar_colors = ["#FF0000", "#00A08A", "#F2AD00", "#F98400", "#5BBCD6", "#ECCBAE", "#046C9A", "#D69C4E", "#ABDDDE", "#000000"]
resource_icons = {".mp4": "fa-video-camera", ".png": "fa-file-image-o", ".pdf": "fa-file-pdf-o", ".zip": "fa-file-archive-o"}

format_duration = lambda t: str(timedelta(seconds=t.seconds))

def get_run_stats(current_run_stats, previous_run_stats, format_value_fn = lambda x: x):
    if not current_run_stats:
        return []
    stats = []
    for k, v in current_run_stats.items():
        prev_value = previous_run_stats.get(k, 0) if previous_run_stats else 0
        bg_class = "table-default"
        if v < prev_value:
            bg_class = "table-danger"
        elif v > prev_value:
            bg_class = "table-success"
        stats.append({
                "icon": resource_icons.get(k, "fa-file"),
                "name": k,
                "value": format_value_fn(v),
                "previous_value": format_value_fn(prev_value) if prev_value else "-",
                "bg_class": bg_class,
            })
    return stats

class RunView(TemplateView):
    template_name = "runs.html"

    @method_decorator(ensure_csrf_cookie) # via https://stackoverflow.com/a/43712324/127114
    def get(self, request, *args, **kwargs):
        return super(RunView, self).get(request, *args, **kwargs)

    def get_context_data(self, **kwargs):
        context = super(RunView, self).get_context_data(**kwargs)
        run_id = uuid.UUID(kwargs.get('runid', ''))
        run = ContentChannelRun.objects.get(run_id=run_id)
        # TODO(arvnd): The previous run will be wrong for any run that
        # is not the most recent.
        previous_run = run.channel.runs.all()[:2]
        if len(previous_run) < 2:
            previous_run = None
        else:
            previous_run = previous_run[1]
        context['channel'] = run.channel
        context['run_stages'] = []
        total_time = timedelta()
        for idx, stage in enumerate(run.events.order_by('finished').all()):
            context['run_stages'].append({
                "duration": stage.duration,
                "name": stage.name.replace("Status.",""),
                "color": progress_bar_colors[idx % len(progress_bar_colors)]})
            total_time += stage.duration
        for stage in context['run_stages']:
            stage['percentage'] = stage['duration'] / total_time * 100
            stage['duration'] = format_duration(stage['duration'])
        context['total_time'] = format_duration(total_time)

        context['resource_counts'] = get_run_stats(run.resource_counts, previous_run.resource_counts)
        context['resource_sizes'] = get_run_stats(run.resource_sizes, previous_run.resource_sizes, sizeof_fmt)

        if self.request.user in run.channel.followers.all():
            # closed star if the user has already saved this.
            context['saved_icon_class'] = 'fa-star'
        else:
            context['saved_icon_class'] = 'fa-star-o'

        # todo graphs
        # todo save to my profile

        return context<commit_msg>small bugfix for runs without timing<commit_after>import uuid
from datetime import timedelta
from datetime import time

from django.conf import settings
from django.utils.decorators import method_decorator
from django.views.decorators.csrf import ensure_csrf_cookie
from django.views.generic.base import TemplateView
from django.utils.safestring import mark_safe

from .models import *

import redis
# TODO(arvnd): Should probably get the redis adapter from some common
# place.
REDIS = redis.StrictRedis(host=settings.MMVP_REDIS_HOST,
                          port=settings.MMVP_REDIS_PORT,
                          db=settings.MMVP_REDIS_DB,
                          charset="utf-8",
                          decode_responses=True)

def sizeof_fmt(num, suffix='B'):
    for unit in ['','K','M','G']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'T', suffix)

# Darjeeling Limited
progress_bar_colors = ["#FF0000", "#00A08A", "#F2AD00", "#F98400", "#5BBCD6", "#ECCBAE", "#046C9A", "#D69C4E", "#ABDDDE", "#000000"]
resource_icons = {".mp4": "fa-video-camera", ".png": "fa-file-image-o", ".pdf": "fa-file-pdf-o", ".zip": "fa-file-archive-o"}

format_duration = lambda t: str(timedelta(seconds=t.seconds))

def get_run_stats(current_run_stats, previous_run_stats, format_value_fn = lambda x: x):
    if not current_run_stats:
        return []
    stats = []
    for k, v in current_run_stats.items():
        prev_value = previous_run_stats.get(k, 0) if previous_run_stats else 0
        bg_class = "table-default"
        if v < prev_value:
            bg_class = "table-danger"
        elif v > prev_value:
            bg_class = "table-success"
        stats.append({
                "icon": resource_icons.get(k, "fa-file"),
                "name": k,
                "value": format_value_fn(v),
                "previous_value": format_value_fn(prev_value) if prev_value else "-",
                "bg_class": bg_class,
            })
    return stats

class RunView(TemplateView):
    template_name = "runs.html"

    @method_decorator(ensure_csrf_cookie) # via https://stackoverflow.com/a/43712324/127114
    def get(self, request, *args, **kwargs):
        return super(RunView, self).get(request, *args, **kwargs)

    def get_context_data(self, **kwargs):
        context = super(RunView, self).get_context_data(**kwargs)
        run_id = uuid.UUID(kwargs.get('runid', ''))
        run = ContentChannelRun.objects.get(run_id=run_id)
        # TODO(arvnd): The previous run will be wrong for any run that
        # is not the most recent.
        previous_run = run.channel.runs.all()[:2]
        if len(previous_run) < 2:
            previous_run = None
        else:
            previous_run = previous_run[1]
        context['channel'] = run.channel
        context['run_stages'] = []
        total_time = timedelta()
        for idx, stage in enumerate(run.events.order_by('finished').all()):
            context['run_stages'].append({
                "duration": stage.duration,
                "name": stage.name.replace("Status.",""),
                "color": progress_bar_colors[idx % len(progress_bar_colors)]})
            total_time += stage.duration
        for stage in context['run_stages']:
            stage['percentage'] = stage['duration'] / total_time * 100 if total_time.seconds > 0 else 0
            stage['duration'] = format_duration(stage['duration'])
        context['total_time'] = format_duration(total_time)

        context['resource_counts'] = get_run_stats(run.resource_counts, previous_run.resource_counts)
        context['resource_sizes'] = get_run_stats(run.resource_sizes, previous_run.resource_sizes, sizeof_fmt)

        if self.request.user in run.channel.followers.all():
            # closed star if the user has already saved this.
            context['saved_icon_class'] = 'fa-star'
        else:
            context['saved_icon_class'] = 'fa-star-o'

        # todo graphs
        # todo save to my profile

        return context<|endoftext|>