import os
import sys
import tempfile
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'payload', 'user', 'remote_access', 'pager-webui'))
import server


def setUpModule():
    __import__('importlib').reload(server)


class HandshakesTest(unittest.TestCase):
    def setUp(self):
        self.dir = tempfile.mkdtemp()
        server.LOOT_HS_DIR = self.dir
        with open(os.path.join(self.dir, 'hs1.cap'), 'wb') as f:
            f.write(b'\x00\x01\x02')
        with open(os.path.join(self.dir, 'hs2.pcap'), 'wb') as f:
            f.write(b'\x00' * 5)

    def test_list(self):
        status, payload = server.h_handshakes_get(type('C', (), {'args': ()})())
        self.assertEqual(status, 200)
        names = [f['name'] for f in payload['files']]
        self.assertEqual(names, ['hs1.cap', 'hs2.pcap'])
        self.assertEqual(payload['files'][0]['size'], 3)

    def test_delete_by_body(self):
        server.h_handshakes_delete(type('C', (), {'args': (), 'body': {'name': 'hs1.cap'}})())
        self.assertFalse(os.path.exists(os.path.join(self.dir, 'hs1.cap')))

    def test_delete_rejects_traversal(self):
        status, payload = server.h_handshakes_delete(type('C', (), {'args': (), 'body': {'name': '../server.py'}})())
        self.assertEqual(status, 400)


class LootProxyTest(unittest.TestCase):
    def test_loot_zip_returns_download(self):
        server.daemon_call = lambda m, p, body=None, token=None, timeout=15: (200, b'PK\x03\x04zipdata')
        server.current_token = lambda: 'tok'
        status, payload = server.h_loot_zip(type('C', (), {'args': ()})())
        self.assertEqual(status, 200)
        self.assertEqual(payload.ctype, 'application/zip')

    def test_loot_archive_posts_to_daemon(self):
        calls = []
        def fake(m, p, body=None, token=None, timeout=15):
            calls.append((m, p))
            return 200, {'ok': True}
        server.daemon_call = fake
        server.current_token = lambda: 'tok'
        server.h_loot_archive(type('C', (), {'args': ()})())
        self.assertTrue(any(m == 'POST' and p == '/api/loot/archive' for m, p in calls))


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