import io
import zipfile
import unittest

import vivisect
import vivisect.tests.helpers as helpers


class MalwareTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        testpath = helpers.getTestPath('windows', 'i386', 'malware.zip')
        cls.embed_vw = vivisect.VivWorkspace()
        cls.mmap_vw = vivisect.VivWorkspace()
        with zipfile.ZipFile(testpath, mode='r') as zp:
            fd = io.BytesIO(zp.open('mal_memmaps.exe', pwd=b'infected', mode='r').read())
            # DEV: intentionally skipping analysis for now since the tests I have in mind
            # initially don't need it (and it adds a lot of time that isn't worth it until we
            # add a lot more tests)
            cls.mmap_vw.loadFromFd(fd)

            # DEV: This one needs analysis, but luckily only takes ~2 seconds to analyze
            fd = io.BytesIO(zp.open('mal_carve.exe', pwd=b'infected', mode='r').read())
            cls.embed_vw.loadFromFd(fd)
            cls.embed_vw.analyze()

    def test_memorymaps_basicperms(self):
        vw = self.mmap_vw
        maps = set([
            (268435456, 0x1000, 4, '0e25aa791c9119108af073bc9e9d0fa2'),
            (268439552, 0x66c00, 5, '0e25aa791c9119108af073bc9e9d0fa2'),
            (268861440, 0xdc00, 4, '0e25aa791c9119108af073bc9e9d0fa2'),
            (268918784, 0x13600, 6, '0e25aa791c9119108af073bc9e9d0fa2'),
            (269004800, 0x4600, 4, '0e25aa791c9119108af073bc9e9d0fa2'),
        ])
        self.assertEqual(maps, set(vw.getMemoryMaps()))

    def test_locations(self):
        # make sure none of the locations in any of the other maps are opcodes
        vw = self.mmap_vw
        for lva, lsize, ltype, linfo in vw.getLocations():
            memmap = vw.getMemoryMap(lva)
            if memmap[0] != 0x10001000:
                self.assertNotEqual(ltype, 5)

    def test_carve_pe(self):
        # all embedded PEs should be marked as dead data
        vw = self.embed_vw
        self.assertTrue(vw.isDeadData(0x10004030))
        self.assertTrue(vw.isDeadData(0x10018074))
        self.assertTrue(vw.isDeadData(0x1001e030))
