96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import struct
|
|
import unittest
|
|
import warnings
|
|
|
|
warnings.filterwarnings("ignore")
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
import Win_Client
|
|
import pgpy
|
|
from pgpy.constants import PubKeyAlgorithm, KeyFlags, HashAlgorithm, SymmetricKeyAlgorithm, CompressionAlgorithm
|
|
|
|
|
|
class TestWinClientComponent(unittest.TestCase):
|
|
def setUp(self):
|
|
self.dummy_config = "test_win_client_config.json"
|
|
# Generate dummy PGP key for testing
|
|
key = pgpy.PGPKey.new(PubKeyAlgorithm.RSAEncryptOrSign, 2048)
|
|
uid = pgpy.PGPUID.new("TestHub")
|
|
key.add_uid(
|
|
uid,
|
|
usage={KeyFlags.EncryptCommunications, KeyFlags.EncryptStorage},
|
|
hashes=[HashAlgorithm.SHA256],
|
|
ciphers=[SymmetricKeyAlgorithm.AES256],
|
|
compression=[CompressionAlgorithm.Uncompressed]
|
|
)
|
|
self.server_priv = key
|
|
self.server_pub = key.pubkey
|
|
self.fingerprint = str(key.pubkey.fingerprint)
|
|
|
|
with open(self.dummy_config, "w", encoding="utf-8") as f:
|
|
json.dump({
|
|
"server_host": "127.0.0.1",
|
|
"server_port": 9443,
|
|
"server_fingerprint": self.fingerprint,
|
|
"server_public_key": str(self.server_pub),
|
|
"auth_token": "secret-test-token"
|
|
}, f)
|
|
|
|
def tearDown(self):
|
|
if os.path.exists(self.dummy_config):
|
|
try:
|
|
os.remove(self.dummy_config)
|
|
except Exception:
|
|
pass
|
|
|
|
def test_client_config_anonymity(self):
|
|
config = Win_Client.load_config(self.dummy_config)
|
|
self.assertNotIn("server_name", config)
|
|
self.assertNotIn("name", config)
|
|
self.assertNotIn("site_name", config)
|
|
self.assertEqual(config["server_fingerprint"], self.fingerprint)
|
|
|
|
def test_get_machine_identifier(self):
|
|
machine_id = Win_Client.get_machine_identifier()
|
|
self.assertIsInstance(machine_id, str)
|
|
self.assertGreater(len(machine_id), 0)
|
|
self.assertNotEqual(machine_id, "localhost")
|
|
|
|
def test_encryption_and_envelope_creation(self):
|
|
config = Win_Client.load_config(self.dummy_config)
|
|
logs = [{
|
|
"server": Win_Client.get_machine_identifier(),
|
|
"signature": "TestWinSignature",
|
|
"severity": "WARNING",
|
|
"message": "Disk space threshold warning"
|
|
}]
|
|
|
|
pub_key, _ = pgpy.PGPKey.from_blob(config["server_public_key"])
|
|
payload = {
|
|
"server": Win_Client.get_machine_identifier(),
|
|
"logs": logs
|
|
}
|
|
msg = pgpy.PGPMessage.new(json.dumps(payload))
|
|
enc = pub_key.encrypt(msg)
|
|
self.assertTrue(str(enc).startswith("-----BEGIN PGP MESSAGE-----"))
|
|
|
|
# Decrypt with private key to verify end-to-end payload integrity
|
|
dec = self.server_priv.decrypt(enc)
|
|
restored = json.loads(dec.message)
|
|
self.assertEqual(restored["logs"][0]["signature"], "TestWinSignature")
|
|
|
|
def test_framing_protocol(self):
|
|
envelope_data = json.dumps({"test": "data"}).encode("utf-8")
|
|
frame = struct.pack(">I", len(envelope_data)) + envelope_data
|
|
self.assertEqual(len(frame), 4 + len(envelope_data))
|
|
length = struct.unpack(">I", frame[:4])[0]
|
|
self.assertEqual(length, len(envelope_data))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|