51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
import os
|
|
import pytest
|
|
from ayn_antivirus.quarantine.vault import QuarantineVault
|
|
|
|
def test_quarantine_and_restore(tmp_path):
|
|
vault_dir = tmp_path / "vault"
|
|
key_file = tmp_path / "keys" / "vault.key"
|
|
vault = QuarantineVault(str(vault_dir), str(key_file))
|
|
|
|
test_file = tmp_path / "malware.txt"
|
|
test_file.write_text("this is malicious content")
|
|
|
|
threat_info = {
|
|
"threat_name": "TestVirus",
|
|
"threat_type": "virus",
|
|
"severity": "high"
|
|
}
|
|
qid = vault.quarantine_file(str(test_file), threat_info)
|
|
assert qid is not None
|
|
assert not test_file.exists()
|
|
assert vault.count() == 1
|
|
|
|
restore_path = tmp_path / "restored.txt"
|
|
vault.restore_file(qid, str(restore_path))
|
|
assert restore_path.exists()
|
|
assert restore_path.read_text() == "this is malicious content"
|
|
|
|
def test_quarantine_list(tmp_path):
|
|
vault_dir = tmp_path / "vault"
|
|
key_file = tmp_path / "keys" / "vault.key"
|
|
vault = QuarantineVault(str(vault_dir), str(key_file))
|
|
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_text("content")
|
|
vault.quarantine_file(str(test_file), {"threat_name": "Test", "threat_type": "virus", "severity": "low"})
|
|
|
|
items = vault.list_quarantined()
|
|
assert len(items) == 1
|
|
|
|
def test_quarantine_delete(tmp_path):
|
|
vault_dir = tmp_path / "vault"
|
|
key_file = tmp_path / "keys" / "vault.key"
|
|
vault = QuarantineVault(str(vault_dir), str(key_file))
|
|
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_text("content")
|
|
qid = vault.quarantine_file(str(test_file), {"threat_name": "Test", "threat_type": "virus", "severity": "low"})
|
|
|
|
assert vault.delete_file(qid) == True
|
|
assert vault.count() == 0
|