import hashlib
import io
import os
import random
from requests import Session, put, get
import shutil
from urllib.parse import urljoin

DEFAULT_REPO_URL = 'http://localhost:5000'
MALWARE_PATH = ''


def requests_download_file(url, sha256, dest_dir):
    r = get(url, stream=True)
    if dest_dir is None:
        b = io.BytesIO()
        shutil.copyfileobj(r.raw, b)
        return b.getvalue()

    with open(os.path.join(dest_dir, sha256), 'wb') as f:
        shutil.copyfileobj(r.raw, f)
    return os.path.join(dest_dir, sha256)


def generate_local_metadata(file_content, malicious):
    sha256 = hashlib.sha256(file_content).hexdigest()
    md5 = hashlib.md5(file_content).hexdigest()
    return {'file_name': sha256, 'sha256': sha256, 'malicious': malicious, 'md5': md5}


class LiveServerSession(Session):
    def __init__(self, prefix_url=None, *args, **kwargs):
        super(LiveServerSession, self).__init__(*args, **kwargs)
        self.prefix_url = prefix_url

    def request(self, method, url, *args, **kwargs):
        url = urljoin(self.prefix_url, url)
        return super(LiveServerSession, self).request(method, url, *args, **kwargs)


# todo define base class with not-implemented


class DummyMalwareRepoClient(object):

    eicar_bytes = b'X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*'
    not_eicar_bytes = b'not a malicious file'

    def __init__(self, client=None):
        pass

    def get_random_file(self, malicious_filter=True):
        if malicious_filter:
            return generate_local_metadata(self.eicar_bytes, malicious_filter), self.eicar_bytes
        return generate_local_metadata(self.not_eicar_bytes, malicious_filter), self.not_eicar_bytes


class PolySwarmMalwareRepoClient(object):

    def __init__(self, client=None):
        self.client = LiveServerSession(prefix_url=os.environ.get('POLYSWARM_MALWARE_REPO_URL', DEFAULT_REPO_URL)) if client is None else client

    def get_file(self, sha256, md_cache=None, dest_dir=None):
        md = self._get_md_via_hash(sha256) if md_cache is None else md_cache
        download_url_signed = md['url']
        file_obj = requests_download_file(download_url_signed, sha256, dest_dir)
        return md, file_obj

    def _get_md_via_hash(self, sha256):
        return self.client.get('/api/artifacts/{0}'.format(sha256)).json

    def is_malicious(self, sha256):
        return self._get_md_via_hash(sha256)['malicious']

    def get_random_file(self, dest_dir=None, malicious_filter=None, seed=None):
        # todo make this serverside when DB gets bigger
        # todo add arbitrary selection for files based on metadata, do serverside

        l = self.client.get('/api/artifacts/').json
        md = l if isinstance(l, list) else l()

        if malicious_filter is not None:
            md = list(filter(lambda x: x.get('malicious', False) == malicious_filter, md))

        if seed:
            # make this deterministic
            random.seed(seed)

        choice = random.choice(md)

        return self.get_file(choice['sha256'], md_cache=choice, dest_dir=dest_dir)

    def push_file(self, content, filename, parent_sha256=None, classifications=['confidential']):

        h = hashlib.sha256(content).hexdigest()
        md5 = hashlib.md5(content).hexdigest()
        size = len(content)
        sha1 = hashlib.sha1(content).hexdigest()
        # todo mimetype
        metadata = {
            'filename': filename,
            'hash': h,
            'md5': md5,
            'size': size,
            'sha1': sha1,
            'classifications': classifications
        }
        if parent_sha256:
            metadata['parent'] = parent_sha256

        where_to_upload = self.client.post('/api/artifacts', json=metadata).json
        url = where_to_upload['url']
        # todo metadata for storage?
        # need another client that's not on the local
        return put(url, data=content)
