import base64
import io
import zipfile

from .common import Common
from .vocabulary import Malware as m
from .vocabulary import ThreatExchange as t


class Malware(Common):

    _URL = t.URL + t.VERSION + t.MALWARE_ANALYSES
    _DETAILS = t.URL + t.VERSION
    _RELATED = t.URL + t.VERSION

    _fields = [
        m.ADDED_ON,
        m.CRX,
        m.ID,
        m.IMPHASH,
        m.METADATA,
        m.MD5,
        m.PASSWORD,
        m.PE_RICH_HEADER,
        m.SAMPLE,
        m.SAMPLE_SIZE,
        m.SAMPLE_SIZE_COMPRESSED,
        m.SAMPLE_TYPE,
        m.SHA1,
        m.SHA256,
        m.SHARE_LEVEL,
        m.SSDEEP,
        m.STATUS,
        m.TAGS,
        m.VICTIM_COUNT,
        m.XPI,
    ]

    _default_fields = [
        m.ADDED_ON,
        m.CRX,
        m.ID,
        m.IMPHASH,
        m.METADATA,
        m.MD5,
        m.PASSWORD,
        m.PE_RICH_HEADER,
        m.SAMPLE_SIZE_COMPRESSED,
        m.SAMPLE_TYPE,
        m.SHA1,
        m.SHA256,
        m.SHARE_LEVEL,
        m.SSDEEP,
        m.STATUS,
        m.TAGS,
        m.VICTIM_COUNT,
        m.XPI,
    ]

    _unique = [
        'zfh',
        'zf',
        'rfh',
        'rf',
    ]

    @property
    def zfh(self):
        """
        Return a file handle of the base64-decoded sample in a zip file.
        """

        if self.get(m.SAMPLE) is None:
            self.details()
        zfh = io.StringIO()
        zfh.write(base64.b64decode(self.get(m.SAMPLE)))
        zfh.seek(0)
        return zfh

    @property
    def zf(self):
        """
        Return the base64-decoded sample in a zip file.
        """

        return self.zfh.read()

    @property
    def rfh(self):
        """
        Return a file handle of the base64-decoded and unzipped sample.
        """

        zfh = self.zfh
        rfh = io.StringIO()
        with zipfile.ZipFile(zfh, 'r') as zf:
            for entry in zf.infolist():
                rfh.write(zf.read(entry.filename,
                                  self.get(m.PASSWORD)))
                rfh.seek(0)
        return rfh

    @property
    def rf(self):
        """
        Return the base64-decoded and unzipped sample.
        """

        return self.rfh.read()
