import json

from .quarantine import *
from maldump.utils import CustomArc4


class Malwarebytes(Quarantine):
    """Implements Malwarebytes quarantine format"""

    def __init__(self):
        super().__init__()
        self.name = 'Malwarebytes'
        self.location = Path('ProgramData/Malwarebytes/MBAMService/Quarantine')

    def _decrypt(self, input):
        key = [
            0x03, 0x7a, 0x55, 0xc5,
            0xdf, 0x39, 0xd2, 0x89,
            0x7f, 0xef, 0x88, 0x3d,
            0xc0, 0x47, 0xb3, 0x17
        ]
        return CustomArc4(key).decode(input)

    def _normalize_time(self, date):
        return dt.strptime(date, '%Y-%m-%dT%H:%M:%SZ')

    def export(self):
        quarfiles = []
        for metafile in self.location.glob('*.data'):
            with open(metafile, 'rb') as f:
                decrypted = self._decrypt(f.read())
            metadata = json.loads(decrypted)

            # Open corresponding malfile
            uid = metadata['trace']['objectID']
            with open(self.location / (uid + '.quar'), 'rb') as f:
                malfile = self._decrypt(f.read())

            q = QuarEntry()
            q.timestamp = self._normalize_time(metadata['trace']['cleanTime'])
            q.threat = metadata['threatName']
            q.path = metadata['trace']['objectPath']
            q.size = len(malfile)
            q.md5 = metadata['trace']['objectMD5'].lower()
            q.malfile = malfile
            quarfiles.append(q)

        return quarfiles
