290 lines
11 KiB
Python
290 lines
11 KiB
Python
import base64
|
|
import contextlib
|
|
import io
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from Crypto.Cipher import AES
|
|
from Crypto.Hash import SHA256
|
|
from Crypto.Protocol.KDF import HKDF
|
|
from Crypto.PublicKey import ECC
|
|
from Crypto.Signature import eddsa
|
|
|
|
from base.secspider import build_secspider_package
|
|
from secspider_tool import main as secspider_main
|
|
|
|
|
|
class TestSecSpiderBuilder(unittest.TestCase):
|
|
def test_build_package_emits_required_headers_and_payload(self):
|
|
private_key = ECC.generate(curve="Ed25519")
|
|
package_text = build_secspider_package(
|
|
source_text="class Spider:\n pass\n",
|
|
name="[直] omofun",
|
|
version="1",
|
|
remark="",
|
|
kid="kid-1",
|
|
signing_private_key=private_key,
|
|
master_secret=b"0123456789abcdef0123456789abcdef",
|
|
)
|
|
|
|
self.assertIn("//@name:[直] omofun", package_text)
|
|
self.assertIn("//@version:1", package_text)
|
|
self.assertIn("//@remark:", package_text)
|
|
self.assertIn("//@format:secspider/1", package_text)
|
|
self.assertIn("//@alg:aes-256-gcm", package_text)
|
|
self.assertIn("//@wrap:hkdf-aes-keywrap", package_text)
|
|
self.assertIn("//@sign:ed25519", package_text)
|
|
self.assertIn("//@kid:kid-1", package_text)
|
|
self.assertIn("//@nonce:base64:", package_text)
|
|
self.assertIn("//@ek:base64:", package_text)
|
|
self.assertIn("//@hash:sha256:", package_text)
|
|
self.assertIn("//@sig:base64:", package_text)
|
|
self.assertIn("payload.base64:", package_text)
|
|
|
|
def test_build_package_signature_verifies_over_headers_and_payload(self):
|
|
private_key = ECC.generate(curve="Ed25519")
|
|
package_text = build_secspider_package(
|
|
source_text="class Spider:\n pass\n",
|
|
name="fixture",
|
|
version="3",
|
|
remark="demo",
|
|
kid="kid-sign",
|
|
signing_private_key=private_key,
|
|
master_secret=b"0123456789abcdef0123456789abcdef",
|
|
)
|
|
|
|
headers = {}
|
|
payload_b64 = ""
|
|
for line in package_text.splitlines():
|
|
if line.startswith("//@"):
|
|
key, _, value = line[3:].partition(":")
|
|
headers[key] = value
|
|
elif line.startswith("payload.base64:"):
|
|
payload_b64 = line.removeprefix("payload.base64:")
|
|
|
|
signing_bytes = "\n".join(
|
|
[
|
|
f"//@name:{headers['name']}",
|
|
f"//@version:{headers['version']}",
|
|
f"//@remark:{headers['remark']}",
|
|
f"//@format:{headers['format']}",
|
|
f"//@alg:{headers['alg']}",
|
|
f"//@wrap:{headers['wrap']}",
|
|
f"//@sign:{headers['sign']}",
|
|
f"//@kid:{headers['kid']}",
|
|
f"//@nonce:{headers['nonce']}",
|
|
f"//@ek:{headers['ek']}",
|
|
f"//@hash:{headers['hash']}",
|
|
f"payload.base64:{payload_b64}",
|
|
]
|
|
).encode("utf-8")
|
|
signature = base64.b64decode(headers["sig"].removeprefix("base64:"))
|
|
|
|
verifier = eddsa.new(private_key.public_key(), "rfc8032")
|
|
verifier.verify(signing_bytes, signature)
|
|
|
|
def test_build_package_encrypts_source_and_wraps_content_key(self):
|
|
private_key = ECC.generate(curve="Ed25519")
|
|
master_secret = b"0123456789abcdef0123456789abcdef"
|
|
source_text = "class Spider:\n value = 'plain'\n"
|
|
package_text = build_secspider_package(
|
|
source_text=source_text,
|
|
name="fixture",
|
|
version="5",
|
|
remark="",
|
|
kid="kid-wrap",
|
|
signing_private_key=private_key,
|
|
master_secret=master_secret,
|
|
)
|
|
|
|
self.assertNotIn(source_text, package_text)
|
|
headers = {}
|
|
payload_b64 = ""
|
|
for line in package_text.splitlines():
|
|
if line.startswith("//@"):
|
|
key, _, value = line[3:].partition(":")
|
|
headers[key] = value
|
|
elif line.startswith("payload.base64:"):
|
|
payload_b64 = line.removeprefix("payload.base64:")
|
|
|
|
wrap_key = HKDF(
|
|
master=master_secret,
|
|
key_len=32,
|
|
salt=headers["kid"].encode("utf-8"),
|
|
hashmod=SHA256,
|
|
num_keys=1,
|
|
context=f"secspider:{headers['name']}:{headers['version']}:wrap-key".encode("utf-8"),
|
|
)
|
|
wrap_nonce = HKDF(
|
|
master=master_secret,
|
|
key_len=12,
|
|
salt=headers["kid"].encode("utf-8"),
|
|
hashmod=SHA256,
|
|
num_keys=1,
|
|
context=f"secspider:{headers['name']}:{headers['version']}:wrap-nonce".encode("utf-8"),
|
|
)
|
|
wrap_blob = base64.b64decode(headers["ek"].removeprefix("base64:"))
|
|
wrap_cipher = AES.new(wrap_key, AES.MODE_GCM, nonce=wrap_nonce)
|
|
content_key = wrap_cipher.decrypt_and_verify(wrap_blob[:-16], wrap_blob[-16:])
|
|
|
|
payload_blob = base64.b64decode(payload_b64)
|
|
payload_nonce = base64.b64decode(headers["nonce"].removeprefix("base64:"))
|
|
payload_cipher = AES.new(content_key, AES.MODE_GCM, nonce=payload_nonce)
|
|
decrypted = payload_cipher.decrypt_and_verify(payload_blob[:-16], payload_blob[-16:]).decode("utf-8")
|
|
|
|
self.assertEqual(decrypted, source_text)
|
|
|
|
|
|
class TestSecSpiderCli(unittest.TestCase):
|
|
def test_genkeys_writes_private_and_public_key_files(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir)
|
|
private_path = temp_path / "signing-private.pem"
|
|
public_path = temp_path / "signing-public.pem"
|
|
stdout = io.StringIO()
|
|
|
|
with contextlib.redirect_stdout(stdout):
|
|
exit_code = secspider_main(
|
|
[
|
|
"genkeys",
|
|
"--private-key",
|
|
str(private_path),
|
|
"--public-key",
|
|
str(public_path),
|
|
]
|
|
)
|
|
|
|
self.assertEqual(exit_code, 0)
|
|
self.assertTrue(private_path.is_file())
|
|
self.assertTrue(public_path.is_file())
|
|
self.assertIn("wrote", stdout.getvalue())
|
|
|
|
def test_pack_builds_secspider_file_from_source_and_key_files(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir)
|
|
private_path = temp_path / "signing-private.pem"
|
|
public_path = temp_path / "signing-public.pem"
|
|
secret_path = temp_path / "master-secret.txt"
|
|
source_path = temp_path / "fixture.py"
|
|
output_path = temp_path / "fixture.sec.py"
|
|
|
|
secspider_main(
|
|
[
|
|
"genkeys",
|
|
"--private-key",
|
|
str(private_path),
|
|
"--public-key",
|
|
str(public_path),
|
|
]
|
|
)
|
|
secret_path.write_text("0123456789abcdef0123456789abcdef", encoding="utf-8")
|
|
source_path.write_text("class Spider:\n pass\n", encoding="utf-8")
|
|
|
|
exit_code = secspider_main(
|
|
[
|
|
"pack",
|
|
"--input",
|
|
str(source_path),
|
|
"--output",
|
|
str(output_path),
|
|
"--name",
|
|
"fixture",
|
|
"--version",
|
|
"9",
|
|
"--remark",
|
|
"",
|
|
"--kid",
|
|
"kid-cli",
|
|
"--private-key",
|
|
str(private_path),
|
|
"--master-secret-file",
|
|
str(secret_path),
|
|
]
|
|
)
|
|
|
|
self.assertEqual(exit_code, 0)
|
|
package_text = output_path.read_text(encoding="utf-8")
|
|
self.assertIn("//@name:fixture", package_text)
|
|
self.assertIn("//@version:9", package_text)
|
|
self.assertIn("//@kid:kid-cli", package_text)
|
|
self.assertIn("payload.base64:", package_text)
|
|
|
|
def test_pack_uses_default_name_version_and_secret_file_paths(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir)
|
|
source_path = temp_path / "红果短剧.py"
|
|
output_path = temp_path / "红果短剧.sec.py"
|
|
private_path = temp_path / "signing-private.pem"
|
|
public_path = temp_path / "signing-public.pem"
|
|
secret_path = temp_path / "master-secret.txt"
|
|
|
|
secspider_main(
|
|
[
|
|
"genkeys",
|
|
"--private-key",
|
|
str(private_path),
|
|
"--public-key",
|
|
str(public_path),
|
|
]
|
|
)
|
|
secret_path.write_text("0123456789abcdef0123456789abcdef", encoding="utf-8")
|
|
source_path.write_text("class Spider:\n pass\n", encoding="utf-8")
|
|
|
|
with contextlib.chdir(temp_path):
|
|
exit_code = secspider_main(
|
|
[
|
|
"pack",
|
|
"--input",
|
|
str(source_path),
|
|
"--output",
|
|
str(output_path),
|
|
]
|
|
)
|
|
|
|
self.assertEqual(exit_code, 0)
|
|
package_text = output_path.read_text(encoding="utf-8")
|
|
self.assertIn("//@name:红果短剧", package_text)
|
|
self.assertIn("//@version:1", package_text)
|
|
self.assertIn("//@remark:", package_text)
|
|
self.assertIn("//@kid:k2026_04", package_text)
|
|
|
|
def test_pack_uses_default_output_name_txt_suffix(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir)
|
|
source_path = temp_path / "默认命名.py"
|
|
private_path = temp_path / "signing-private.pem"
|
|
public_path = temp_path / "signing-public.pem"
|
|
secret_path = temp_path / "master-secret.txt"
|
|
|
|
secspider_main(
|
|
[
|
|
"genkeys",
|
|
"--private-key",
|
|
str(private_path),
|
|
"--public-key",
|
|
str(public_path),
|
|
]
|
|
)
|
|
secret_path.write_text("0123456789abcdef0123456789abcdef", encoding="utf-8")
|
|
source_path.write_text("class Spider:\n pass\n", encoding="utf-8")
|
|
|
|
with contextlib.chdir(temp_path):
|
|
exit_code = secspider_main(
|
|
[
|
|
"pack",
|
|
"--input",
|
|
str(source_path),
|
|
]
|
|
)
|
|
|
|
self.assertEqual(exit_code, 0)
|
|
output_path = temp_path / "默认命名.txt"
|
|
self.assertTrue(output_path.is_file())
|
|
package_text = output_path.read_text(encoding="utf-8")
|
|
self.assertIn("//@name:默认命名", package_text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|