#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import datetime
import io
import json
import logging
import os
import re
import typing

import pyzipper
import requests
import requests.adapters

__version__ = '1.0.0'


class FixedTimeoutAdapter(requests.adapters.HTTPAdapter):
    def send(self, *pargs, **kwargs):
        if kwargs['timeout'] is None:
            kwargs['timeout'] = 10
        return super(FixedTimeoutAdapter, self).send(*pargs, **kwargs)


class DateTimeFactory:
    @staticmethod
    def to_str(dt: datetime) -> str:
        return dt.strftime('%Y-%m-%d %H:%M:%S')

    @staticmethod
    def to_date_str(dt: datetime) -> str:
        return dt.strftime('%Y-%m-%d')

    @staticmethod
    def to_time_str(dt: datetime) -> str:
        return dt.strftime('%H:%M:%S')

    @staticmethod
    def from_str(s: str) -> datetime:
        return datetime.datetime.strptime(s, '%Y-%m-%d %H:%M:%S')


class MalwareBazaarException(Exception):
    pass


class MalwareSample:
    def __init__(
            self, sha256: str, first_seen: datetime, last_seen: datetime, file_name: str, file_size: int,
            file_type: str, signature: str, imphash: str, tlsh: str, ssdeep: str, tags: typing.List[str]
    ):
        self.sha256 = sha256
        self.first_seen = first_seen
        self.last_seen = last_seen
        self.file_name = file_name
        self.file_size = file_size
        self.file_type = file_type
        self.signature = signature
        self.imphash = imphash
        self.tlsh = tlsh
        self.ssdeep = ssdeep
        self.tags = tags

    @staticmethod
    def from_row(row):
        return MalwareSample(
            row['sha256_hash'],
            DateTimeFactory.from_str(row['first_seen']),
            None if row['last_seen'] is None else DateTimeFactory.from_str(row['last_seen']),
            row['file_name'],
            row['file_size'],
            row['file_type'],
            row['signature'],
            row['imphash'],
            row['tlsh'],
            row['ssdeep'],
            row['tags'],
        )

    def __repr__(self):
        return F'<Malware {self.sha256} ' \
               F'signature={self.signature} ' \
               F'first={DateTimeFactory.to_date_str(self.first_seen)} ' \
               F'tags={",".join(self.tags)}>'


class MalwareBazaarClient:
    def __init__(self, api_key: str, user_agent: str, base_url: str):
        self.base_url = base_url

        self.session = requests.session()
        self.session.mount('https://', FixedTimeoutAdapter())
        self.session.mount('http://', FixedTimeoutAdapter())
        self.session.headers = {
            'API-KEY': api_key,
            'User-Agent': user_agent,
        }

    def download(self, sha256):
        response = self.session.post(self.base_url, data={
            'query': 'get_file',
            'sha256_hash': sha256,
        })
        response.raise_for_status()

        data = io.BytesIO()
        data.write(response.content)
        zip_file = pyzipper.AESZipFile(data)
        zip_file.setpassword(b'infected')
        for name in zip_file.namelist():
            return zip_file.read(name)

    def by_signature(self, signature_name: str, limit: int = 10) -> typing.Iterable[MalwareSample]:
        response = self.session.post(self.base_url, data={
            'query': 'get_siginfo',
            'signature': signature_name,
            'limit': limit
        })
        response.raise_for_status()
        j = response.json()
        if j['query_status'] != 'ok':
            raise MalwareBazaarException(F'Invalid Query Status: {j["query_status"]}')
        for row in j['data']:
            yield MalwareSample.from_row(row)


class ConsoleHandler(logging.Handler):
    def emit(self, record):
        print('[%s] %s' % (record.levelname, record.msg))


def main():
    import platform

    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='command')

    download_parser = subparsers.add_parser(
        'download', help='Download a sample / list of samples'
    )
    download_parser.add_argument('sha256s', nargs='+')

    signature_parser = subparsers.add_parser(
        'signature', help='List malware samples by signature: AgentTesla, Trickbot...'
    )
    signature_parser.add_argument('signature_name')
    signature_parser.add_argument('--limit', '-l', type=int, default=10)

    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--json', action='store_true')
    parser.add_argument('--api-key', default=os.environ.get('MALWARE_BAZAAR_API_KEY'))
    parser.add_argument('--base-url', default='https://mb-api.abuse.ch/api/v1')
    parser.add_argument(
        '--user-agent',
        default=F'MalwareBazaarClient/{__version__} (python-requests {requests.__version__}) '
                F'{platform.system()} ({platform.release()})'
    )
    args = parser.parse_args()

    logger = logging.getLogger('MalwareBazaar')
    logger.handlers.append(ConsoleHandler())
    logger.setLevel(logging.DEBUG if args.debug else logging.INFO)
    if args.debug:
        import http.client as http_client
        http_client.HTTPConnection.debuglevel = 1

    logger.debug(F'Using User-Agent string: {args.user_agent}')
    client = MalwareBazaarClient(args.api_key, args.user_agent, args.base_url)
    try:
        if args.command == 'signature':
            for sample in client.by_signature(args.signature_name, args.limit):
                if args.json:
                    print(json.dumps({
                        'sha256': sample.sha256,
                        'tags': sample.tags,
                    }))
                else:
                    print(sample)
        if args.command == 'download':
            skipped = 0
            downloaded = 0
            sha256_r = re.compile(r'^[\da-f]{64}$')
            for sha256 in args.sha256s:
                sha256 = sha256.lower()
                if os.path.exists(sha256):
                    logger.warning(F'File with name {sha256} already exists.')
                    skipped += 1
                    continue
                if not sha256_r.match(sha256):
                    logger.warning(F'Argument "{sha256}" not a valid SHA256 hash.')
                    continue
                with open(sha256, 'wb') as fp:
                    fp.write(client.download(sha256))
                downloaded += 1
            logger.info(F'Downloaded {downloaded} file(s), {skipped} skipped.')

    except MalwareBazaarException as e:
        logger.exception(e)


if __name__ == '__main__':
    main()
