import argparse
import json
import os
import sys
import dpkt
import requests

from .pcap_extractor.MyEncoder import MyEncoder
from .pcap_extractor.FlowExtractor import FlowExtractor
from .pcap_extractor.TCPFlow import TCPFlow
from .pcap_extractor.mail import Mail
from .pcap_extractor.HTTPParser import HTTPParser, HttpData
from .pcap_extractor.SMTPParser import SMTPParser
from .pcap_extractor.POP3Parser import POP3Parser
from .pcap_extractor.IMAPParser import IMAPParser
from .pcap_extractor.DNSExtractor import DNSExtractor, DNSItem
from .report import Report, FileHash
from progress.bar import Bar


class CCTXPcapAnalyser:
    def __init__(self, args):
        self.host = args.host
        self.path = args.path
        self.port = args.port
        self.https = args.https
        self.username = args.username
        self.password = args.password
        self.inputFile = args.pcapfile
        self.outputFile = args.outputfile
        self.progress = args.progress
        if self.port == -1:
            # 不指定端口，使用默认的80或者443
            if self.https:
                self.port = 443
            else:
                self.port = 80

        self.flowExtractor = FlowExtractor(valueCallback=self.dealStream)
        self.dnsExtractor = DNSExtractor(valueCallback=self.dealDNSRecord)
        self.httpParser = HTTPParser()
        self.smtpParser = SMTPParser()
        self.pop3Parser = POP3Parser()
        self.imapParser = IMAPParser()

        self.report = Report()
        self.url = f'{"https" if self.https else "http"}://{self.host}:{self.port}{self.path}'
        self.registerUrl = f'{self.url}/user/register'
        self.loginUrl = f'{self.url}/user/login'
        self.queryUrl = f'{self.url}/user/queryObservables'
        self.token = ""
        self.refreshToken = ""

    def doAuthGet(self, url: str, params: dict):
        r = requests.get(url, params=params, headers={
            "Authorization": f"Bearer {self.token}"
        })
        responseData = r.json()
        if responseData == 401:
            # 如果权限验证过期，重新登录
            self.doLogin()
            return self.doAuthGet(url, params)
        return responseData

    def doAuthPost(self, url: str, jsonData: dict):
        r = requests.post(url, json=jsonData, headers={
            "Authorization": f"Bearer {self.token}"
        })
        responseData = r.json()
        if responseData == 401:
            # 如果权限验证过期，重新登录
            self.doLogin()
            return self.doAuthPost(url, jsonData)
        return responseData

    @staticmethod
    def buildQueryDomain(domain: str) -> dict:
        return {
            'queryType': 'domain',
            'value': domain
        }

    @staticmethod
    def buildQueryAddress(address: str) -> dict:
        return {
            'queryType': 'address',
            'value': address
        }

    @staticmethod
    def buildQueryUri(uri: str) -> dict:
        return {
            'queryType': 'uri',
            'value': uri
        }

    @staticmethod
    def buildQueryFileHash(fileHash: FileHash) -> dict:
        return {
            'queryType': 'fileHash',
            'md5': fileHash.md5,
            'sha1': fileHash.sha1,
            'sha256': fileHash.sha256
        }

    def doQuery(self, queryItems: [dict]) -> [dict]:
        response = self.doAuthGet(self.queryUrl, params={
            'queries': json.dumps(queryItems),
        })
        if response['code'] == 200:
            if response['data'] and len(response['data']) > 0:
                return response['data']
        else:
            print(f'Query Observable failed, {response["msg"]}')
        return []

    def dealDNSRecord(self, dnsRecord: DNSItem):
        """
        处理每个提取到的 DNS 解析记录
        :param dnsRecord:
        :return:
        """
        # query DNSItem
        observables = self.doQuery([
            CCTXPcapAnalyser.buildQueryDomain(dnsRecord.domain),
            CCTXPcapAnalyser.buildQueryAddress(dnsRecord.value)
        ])
        self.report.addDNSRecord(dnsRecord, observables)

    def dealMail(self, mail: Mail, tcpFlow: TCPFlow):
        # query email address, ip address, file hash
        queryTask = [
            CCTXPcapAnalyser.buildQueryAddress(tcpFlow.srcIP),
            CCTXPcapAnalyser.buildQueryAddress(tcpFlow.dstIP),
            CCTXPcapAnalyser.buildQueryAddress(mail.From),
            CCTXPcapAnalyser.buildQueryAddress(mail.To),
        ]
        for mailFile in mail.files:
            fileHash = FileHash(mailFile.fileData)
            queryTask.append(CCTXPcapAnalyser.buildQueryFileHash(fileHash))
        observables = self.doQuery(queryTask)
        self.report.addEmail(mail, tcpFlow, observables)

    def dealHttpData(self, httpData: HttpData, tcpFlow: TCPFlow):
        # query url, file hash, ip address, domain
        queryTask = [
            CCTXPcapAnalyser.buildQueryAddress(httpData.getDomain()),
            CCTXPcapAnalyser.buildQueryUri(httpData.getUrl()),
            CCTXPcapAnalyser.buildQueryAddress(tcpFlow.srcIP),
            CCTXPcapAnalyser.buildQueryAddress(tcpFlow.dstIP),
        ]
        if httpData.fileHashes:
            fileHash = FileHash()
            fileHash.md5 = httpData.fileHashes["md5"]
            fileHash.sha1 = httpData.fileHashes["sha1"]
            fileHash.sha256 = httpData.fileHashes["sha256"]
            queryTask.append(CCTXPcapAnalyser.buildQueryFileHash(fileHash))
        observables = self.doQuery(queryTask)
        self.report.addHttp(httpData, tcpFlow, observables)

    def dealFTP(self, data: bytes, tcpFlow: TCPFlow):
        fileHash = FileHash(data)
        queryTask = [
            CCTXPcapAnalyser.buildQueryFileHash(fileHash),
            CCTXPcapAnalyser.buildQueryAddress(tcpFlow.srcIP),
            CCTXPcapAnalyser.buildQueryAddress(tcpFlow.dstIP),
        ]
        observables = self.doQuery(queryTask)
        self.report.addFTP(fileHash, tcpFlow, observables)

    def dealStream(self, tcpFlow: TCPFlow):
        """
        处理每个提取到的TCP流
        :param tcpFlow:
        :return:
        """
        self.report.addTCPFlow(tcpFlow)
        forwardBytes, reverseBytes = tcpFlow.getAllForwardBytes(), tcpFlow.getAllReverseBytes()
        if tcpFlow.dstPort == 21:
            pass
        elif tcpFlow.srcPort == 20 or tcpFlow.dstPort == 20:
            # 处理 FTP
            data1, data2 = forwardBytes, reverseBytes
            data = data1 if len(data2) == 0 else data2
            if len(data) > 0:
                # md1 = hashlib.md5()
                # md2 = hashlib.md5()
                # md3 = hashlib.md5()
                # with closing(BytesIO(data)) as data:
                #     for line in data.readlines():
                #         md1.update(line)
                #         if line.endswith(b"\r\n"):
                #             md2.update(line[:-2])
                #             md2.update(b'\r')
                #             md3.update(line[:-2])
                #             md3.update(b'\n')
                self.dealFTP(data, tcpFlow)
        elif tcpFlow.dstPort == 143:
            # 处理 IMAP
            for mail in self.imapParser.parse(forwardBytes, reverseBytes):
                self.dealMail(mail, tcpFlow)
        elif tcpFlow.dstPort == 110:
            for mail in self.pop3Parser.parse(reverseBytes):
                self.dealMail(mail, tcpFlow)
        elif tcpFlow.dstPort == 25:
            for mail in self.smtpParser.parse(forwardBytes):
                self.dealMail(mail, tcpFlow)
        elif (len(forwardBytes) == 0 and len(reverseBytes) > 0) or (len(forwardBytes) > 0 and len(reverseBytes) == 0):
            # try to cal file hash for FTP passive mode
            if len(forwardBytes) == 0 and len(reverseBytes) > 0:
                self.dealFTP(reverseBytes, tcpFlow)
            else:
                self.dealFTP(forwardBytes, tcpFlow)
        else:
            # parse http
            for httpData in self.httpParser.parse(forwardBytes, reverseBytes):
                self.dealHttpData(httpData, tcpFlow)

    def doRegister(self):
        r = requests.post(self.registerUrl, json={
            'username': self.username,
            'password': self.password
        })
        return r.json()

    def doLogin(self):
        r = requests.post(self.loginUrl, json={
            'username': self.username,
            'password': self.password
        })
        data = r.json()
        if data['code'] != 200:
            print(f"Login failed, {data['msg']}")
            return False
        self.token = data['data']['token']
        self.refreshToken = data['data']['refresh_token']
        return True

    def start(self):
        """
        Start to parse pcap file
        :return:
        """

        # do login first
        if not self.doLogin():
            return
        # dnsRecord = DNSItem(['qjm253.cn', 'A', '192.168.1.1', 1, 1])
        # dnsRecord.domain = 'qjm253.cn'
        # dnsRecord.value = '192.168.1.1'
        # self.dealDNSRecord(dnsRecord)
        # return

        if self.progress:
            bar = Bar('Extract Progress:', max=100)
        else:
            bar = None
        with open(self.inputFile, 'rb') as pcap:
            progress = 0
            totalSize = os.path.getsize(self.inputFile)
            packets = dpkt.pcap.Reader(pcap)
            for ts, buf in packets:
                if self.progress:
                    currentProgress = int(pcap.tell() * 100 / totalSize)
                    if currentProgress != progress:
                        progress = currentProgress
                        bar.next()
                ethPacket = dpkt.ethernet.Ethernet(buf)
                self.report.addPacket(ethPacket, ts)
                self.flowExtractor.addPacket(ethPacket, ts)
                self.dnsExtractor.addPacket(ethPacket, ts)
        if self.progress:
            bar.finish()
        self.flowExtractor.done()


def main():
    """
    cctxpa is a command lien tool for CCTX to parse pcap file and compare with CCTX's Observables
    """
    if len(sys.argv) == 1:
        sys.argv.append('--help')
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', type=str, default="127.0.0.1", help="CCTX pcap analyser server addresss")
    parser.add_argument('--port', type=int, default=5000, help="CCTX pcap analyser server port")
    parser.add_argument('--https', action='store_true', help="Use https or http")
    parser.add_argument('--register', action='store_true', help="Just do register account")
    parser.add_argument('--path', type=str, default="", help="CCTX pcap analyser server login path")
    parser.add_argument('--progress', action='store_true',
                        help="Print progress, if open, maybe lead slow extract speed.")
    parser.add_argument('-u', '--username', type=str, required=True, help="Username")
    parser.add_argument('-p', '--password', type=str, required=True, help="Password")
    parser.add_argument('-f', '--pcapfile', type=str, default='test.pcap', help="Pcap file need to parse!")
    parser.add_argument('-o', '--outputfile', type=str, default="report.json", help="A file to store output report")
    args = parser.parse_args()

    cctxpa = CCTXPcapAnalyser(args)
    if args.register:
        # do register
        cctxpa.doRegister()
    else:
        cctxpa.start()
        with open(args.outputfile, 'w') as file:
            file.write(json.dumps(cctxpa.report.toDict(), ensure_ascii=False, cls=MyEncoder))


if __name__ == '__main__':
    main()
