test_malware.py

import io
import os
import pathlib
import subprocess
import tarfile
import tempfile
import unittest
from unittest.mock import patch

import malware


class SnapshotTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.root = pathlib.Path(self.temp.name)
        self.git('init', '--initial-branch=master')
        self.git('config', 'user.name', 'Test Author')
        self.git('config', 'user.email', 'author@example.invalid')

    def git(self, *args):
        return subprocess.check_output(['git', '-C', str(self.root), *args], text=True, stderr=subprocess.DEVNULL).strip()

    def commit(self):
        self.git('add', '.')
        self.git('commit', '-m', 'Snapshot test')
        return self.git('rev-parse', 'HEAD')

    def test_snapshot_includes_files_hidden_from_git_archive(self):
        (self.root / '.gitattributes').write_text('hidden.txt export-ignore\n')
        (self.root / 'hidden.txt').write_text('must still be scanned')
        payload, count = malware.snapshot(self.root, self.commit())
        self.assertEqual(count, 2)
        with tarfile.open(fileobj=io.BytesIO(payload)) as archive:
            contents = [archive.extractfile(member).read() for member in archive.getmembers()]
        self.assertIn(b'must still be scanned', contents)

    def test_symlink_is_scanned_as_text_without_following_it(self):
        (self.root / 'link').symlink_to('/etc/passwd')
        payload, _ = malware.snapshot(self.root, self.commit())
        with tarfile.open(fileobj=io.BytesIO(payload)) as archive:
            member = archive.getmembers()[0]
            self.assertTrue(member.isfile())
            self.assertEqual(archive.extractfile(member).read(), b'/etc/passwd')

    def test_lfs_payloads_fail_closed(self):
        (self.root / 'large').write_text('version https://git-lfs.github.com/spec/v1\noid sha256:abc\nsize 999\n')
        with self.assertRaisesRegex(RuntimeError, 'LFS payloads'):
            malware.snapshot(self.root, self.commit())

    def test_oversized_snapshot_fails_before_scanning(self):
        (self.root / 'large').write_text('x' * 50)
        commit = self.commit()
        with patch.object(malware, 'MAX_BYTES', 10):
            with self.assertRaisesRegex(RuntimeError, 'size limit'):
                malware.snapshot(self.root, commit)


class ScannerTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        env = patch.dict(os.environ, {'XDG_STATE_HOME': self.temp.name})
        env.start(); self.addCleanup(env.stop)

    def scan(self, code):
        with patch.object(malware, 'binary', return_value=('/trusted/clamscan', {})), \
             patch.object(malware, 'refresh_signatures'), patch.object(malware, 'snapshot', return_value=(b'pinned snapshot', 4)), \
             patch.object(malware.subprocess, 'run', return_value=subprocess.CompletedProcess([], code)) as scan:
            result = malware.scan_commit('/checkout', 'b' * 40)
            argv = scan.call_args_list[0].args[0]
            self.assertIn('--alert-exceeds-max=yes', argv)
            self.assertIn('--alert-encrypted=yes', argv)
            self.assertEqual(scan.call_args_list[0].kwargs['input'], b'pinned snapshot')
            return result

    def test_clean_scan_returns_receipt_for_the_exact_commit(self):
        self.assertEqual(self.scan(0)['commit'], 'b' * 40)

    def test_detection_blocks_update(self):
        with self.assertRaisesRegex(RuntimeError, 'detected a threat'):
            self.scan(1)

    def test_scanner_error_blocks_update(self):
        with self.assertRaisesRegex(RuntimeError, 'could not complete'):
            self.scan(2)


if __name__ == '__main__':
    unittest.main()