import os
import json
import IPy
from multiprocessing import Pool



def loadFromFile(filename):
    result = []
    for line in open(filename,encoding='latin1').readlines():
        data = json.loads(line)
        flags = data['DNS']['Flags']
        isAnswer = bin(flags)[2]=='1'
        if 'Queries' not in data['DNS']:
            continue
        query = data['DNS']['Queries'][0]
        domain = query['Name']
        if isAnswer and 'Answers' in data['DNS'] and data['DNS']['Answers'] is not None:
            for answer in data['DNS']['Answers']:
                if 'Value' not in answer:
                    continue
                adderss = answer['Value']                
                if answer['Type']==28:
                    # ipType = IPy.IP(adderss).version()
                    # if ipType==4:
                    #     continue
                    result.append({
                        'domain':domain,
                        'address':adderss,
                        'time':data['PacketInfor']['TimeStamp'],
                        'key':domain+'_to_'+adderss
                    })
    print(filename,len(result))
    return result

def loadFiles(feature):
    result = []
    basedir = '/Tsinghua_LanProbe/FlowLog/Extract/'
    for (root,sub,files) in os.walk(basedir):
        for name in files:
            if 'json' not in name:
                continue
            filename = os.path.join(root,name)
            if feature in filename:
                result.append(filename)
    return result


if __name__=='__main__':
    # loadFromFile('/Tsinghua_LanProbe/FlowLog/Extract/1/DNS/2018-08-01_13-13-59_1532958985.json')
    dnsfiles = loadFiles('DNS')
    pool = Pool(20)
    dnsLogs = pool.map(loadFromFile,dnsfiles)
    pool.close()
    pool.join()
    result = {}
    for logs in dnsLogs:
        for log in logs:
            key = log['key'] 
            if key not in result:
                result[key] = {
                    'startTime':log['time'],
                    'endTime':log['time'],
                    'domain':log['domain'],
                    'address':log['address']
                }
            else:
                result[key]['startTime'] = min(result[key]['startTime'] ,log['time'])
                result[key]['endTime'] = max(result[key]['endTime'] ,log['time'])
    json.dump(result,open('ipv6.json','w+'))
