import re
import os
import logging
import time
from malware import utils
from malware.snort import SnortRule
from malware import pickle_tool
from malware import virustotal

logger = logging.getLogger(__name__)

request_times = 0
init_timer = time.time()


def list_pcap(path):
    for dirPath, dirNames, fileNames in os.walk(path):
        for f in fileNames:
            if f.split('.')[1] == 'pcap':
                yield os.path.join(dirPath, f)


def check_timer_counter():
    global request_times
    global init_timer
    if request_times == 300:
        request_times = 0
        period = time.time() - init_timer
        waiting = 60 - period + 1
        if waiting > 0:
            logger.info("Waiting %s seconds", (str(waiting)))
            time.sleep(waiting)
        init_timer = time.time()


def get_positive(resource, cache={}):
    if resource in cache.keys():
        logger.info("%s in cache" % resource)
        positives = cache.get(resource)
        # logger.info('%s in cache positives: %s' % (resource, positives))
        return positives
    else:
        global request_times
        global init_timer
        request_times += 1
        logger.info("Search on VirusTotal counter: %s", str(request_times))
        positives = virustotal.url_report(resource)
        check_timer_counter()

        if positives >= 0:
            cache[resource] = positives
            return positives
        elif positives is None:
            logger.info("No report. Submmit the URL to VirusTotal")
            return None
        else:
            logger.debug("Get reports failed.")
            return None


def iterpayload(path):
    connection = utils.follow_tcp_stream(path)
    for conn, frame in connection.iteritems():
        for seq, content in frame.iteritems():
            if content:
                # Generate the content and 5-tuple
                yield content, conn
            else:
                # Some packets have no payload
                pass


def gen_rule(pcap_path):
    # rule_set = list()
    cache = pickle_tool.check_json()
    # sid = 1000000
    for content, conn in iterpayload('%s' % (pcap_path)):
        # print content, utils.connection_key_2_str(conn)
        get_method = re.search('GET (.*) ', content)
        host = re.search('Host: (.*)', content)
        if host and get_method:
            # pos = get_positive(host.group(1).rstrip(), cache)
            # print host.group(1).rstrip()

            # if pos > 0:
            #     sid += 1
            #     rule = make_host_rule(host.group(0).rstrip(), conn[3])
            #     if not iscontent(rule, RULE):
            #        RULE.append(rule)
                # else:
                #     pass
                    # sid -= 1
            # else:
            #     pass
                # print "NOT generate"
                # NOT generate rule

            if get_method.group(1) == '/':
                url = host.group(1).rstrip()
            else:
                url = host.group(1).rstrip() + get_method.group(1)

            host_content = host.group(0).rstrip()
            uricontent = get_method.group(1)
            pos = get_positive(url, cache)
            if pos > 0:
                # print "BAD URL: %s" % uricontent
                # sid += 1
                if uricontent == '/':
                    uricontent = None
                rule = make_uri_rule(host_content, uricontent, conn[3])
                yield rule
                # if not isuricontent(rule, RULE):
                #    RULE.append(rule)
                # else:
                #    sid -= 1
            else:
                # positives == 0 or positives == None
                pass

        else:
            # print "no URL"
            pass
            # print repr(content)
    # for r in rule_set:
    #     print r
    pickle_tool.update_json(cache)
    # return rule_set


def make_host_rule(content, dst_port, sid=0):
    rule = SnortRule()
    pattern = dict()
    pattern['msg'] = '"Trojan.Gen"'
    pattern['content'] = ['"{host}"'.format(host=content), 'nocase']
    # pattern['sid'] = sid
    pattern['dst_port'] = dst_port
    rule.set_malicious_pattern(**pattern)
    return rule


def make_uri_rule(content, uricontent, dst_port, sid=0):
    rule = SnortRule()
    pattern = dict()
    pattern['msg'] = '"Trojan.Gen.uricontent"'
    pattern['content'] = ['"{host}"'.format(host=content), 'nocase']
    pattern['uricontent'] = ['"{uri}"'.format(uri=uricontent), 'nocase']
    # pattern['sid'] = sid
    pattern['dst_port'] = dst_port
    rule.set_malicious_pattern(**pattern)
    return rule


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

    rule_instances = list()
    RULE = list()

    for pcap in list_pcap('./PCAPLog/'):
        for rule in gen_rule(pcap):
            rule_instances.append(rule)

    for ruleobj in rule_instances:
        RULE.append(str(ruleobj))

    RULE = list(set(RULE))

    for r in RULE:
        print r


if __name__ == "__main__":
    main()

@@ -154,22 +154,72 @@
  import 
-virust
+apikey%0Afrom virus_total_apis import PrivateApi as VirusT
 otal%0A%0Alo
@@ -498,16 +498,17 @@
 )%0A%0A%0Adef 
+_
 check_ti
@@ -857,24 +857,28 @@
 )%0A%0A%0Adef get_
+url_
 positive(res
@@ -927,32 +927,34 @@
 .keys():%0A       
+ #
  logger.info(%22%25s
@@ -1280,24 +1280,65 @@
 est_times))%0A
+        vt = VirusTotal(apikey.APIKEY_0)%0A
         posi
@@ -1350,18 +1350,14 @@
  = v
-irustotal.
+t.get_
 url_
@@ -1364,32 +1364,64 @@
 report(resource)
+.get('results').get('positives')
 %0A        check_t
@@ -1413,16 +1413,17 @@
         
+_
 check_ti
@@ -1633,16 +1633,85 @@
 Total%22)%0A
+            vt.scan_url(resource)%0A            _check_timer_counter()%0A
         
@@ -3337,32 +3337,36 @@
       pos = get_
+url_
 positive(url, ca
@@ -3449,35 +3449,8 @@
 ent%0A
-                # sid += 1%0A
     
@@ -3613,46 +3613,19 @@
     
-yield rule%0A                # if not is
+with open('
 uric
@@ -3626,36 +3626,43 @@
 ('uricontent
-(
+.
 rule
-, RULE)
+s', 'a') as fp
 :%0A          
@@ -3671,30 +3671,49 @@
     
-#
     
-RULE.append
+fp.write('%7Br%7D%5Cn'.format(r=str
 (rule)
+))
 %0A   
@@ -3729,45 +3729,18 @@
     
-# else:%0A                #    sid -= 1
+yield rule
 %0A   

save rules and add virustotal-api module