test: 💍 add unit test

✅ Closes: #16

Author: ZhaoQi99

dev-requirement.txt

@@ -1,3 +1,5 @@
 -r requirement.txt
 isort==5.10.1
-yapf==0.31.0
+yapf==0.31.0
+pytest==6.2.5
+ipython==7.16.3

setup.cfg

@@ -20,7 +20,7 @@ classifiers =
     Programming Language :: Python :: 3.9
     Programming Language :: Python :: Implementation :: CPython
     Environment :: Console
-license_file = LICENSE
+license_files = LICENSE
 description = A simple image converter for Python.
 long_description = file: README.md
 long_description_content_type = text/markdown
@@ -45,4 +45,11 @@ exclude =
 
 [options.entry_points]
 console_scripts =
-    pyencrypt = pyencrypt.cli:cli
+    pyencrypt = pyencrypt.cli:cli
+
+[yapf]
+based_on_style = pep8
+COLUMN_LIMIT = 119
+indent_width =  4
+coalesce_brackets = True
+dedent_closing_brackets = True

tests/constants.py

@@ -0,0 +1 @@
+AES_KEY = b'tiovaZzCF/Hlx/8uw0PXCkiqxCXGosP/AKo6Kn1gECw='

tests/test_aes.py

@@ -0,0 +1,65 @@
+from pyencrypt.aes import AESModeOfOperationECB, add_padding, aes_encrypt, aes_decrypt
+import pytest
+
+from constants import AES_KEY
+
+PLAIN_1 = b'hello world'
+CIPHER_1 = b'\xc5\xa1\xf8\xed\xf7\xa0\x03\xd8\xffu\x01\xac\x93\xcd+\xe1'
+PLAIN_PADDING_1 = PLAIN_1 + b'\x05' * 5
+
+PLAIN_2 = '你好 世界!'.encode()
+CIPHER_2 = b'\t\xb6R0B\x1fgz\x06x\x9d\xaf\xb4\xe7_\x9f'
+PLAIN_PADDING_2 = PLAIN_2 + b'\x02' * 2
+
+PLAIN_3 = b'abcdefghijklmnop'
+CIPHER_3 = b'r\x14\xa7\x92\xd6\x1f\x0c\xf4\x10g\x99\t/\xf0z\xfc' + b'&\x80\xdb\x94\xd1\xf7\x9f\xe0Qo\x05\x98\x7f\xe6j\x8c'
+
+class TestAES:
+
+    @pytest.mark.parametrize('key, plain, cipher', [
+        (AES_KEY, PLAIN_1, CIPHER_1),
+        (AES_KEY, PLAIN_2, CIPHER_2),
+        (AES_KEY, PLAIN_3, CIPHER_3),
+    ])
+    def test_aes_encrypt(self, plain, cipher, key):
+        assert aes_encrypt(plain, key) == cipher
+
+    @pytest.mark.parametrize('key, plain, cipher', [
+        (AES_KEY, PLAIN_1, CIPHER_1),
+        (AES_KEY, PLAIN_2, CIPHER_2),
+        (AES_KEY, PLAIN_3, CIPHER_3),
+    ])
+    def test_aes_decrypt(self, plain, cipher, key):
+        assert aes_decrypt(cipher, key) == plain
+
+
+class TestAESModeOfOperationECB:
+
+    def setup_class(self):
+        self.cipher = AESModeOfOperationECB(AES_KEY)
+
+    @pytest.mark.parametrize('length', [17, 18, 19, 20])
+    def test_encrypt_exception(self, length):
+        with pytest.raises(ValueError) as excinfo:
+            self.cipher.encrypt(b'a' * length)
+        assert str(excinfo.value) == 'plain block must be 16 bytes'
+
+    @pytest.mark.parametrize('length', [17, 18, 19, 20])
+    def test_decrypt_exception(self, length):
+        with pytest.raises(ValueError) as excinfo:
+            self.cipher.decrypt(b'a' * length)
+        assert str(excinfo.value) == 'cipher block must be 16 bytes'
+
+    @pytest.mark.parametrize('plain, cipher', [
+        (PLAIN_PADDING_1, CIPHER_1),
+        (PLAIN_PADDING_2, CIPHER_2),
+    ])
+    def test_encrypt(self,plain, cipher):
+        assert self.cipher.encrypt(plain) == cipher
+
+    @pytest.mark.parametrize('plain, cipher', [
+        (PLAIN_PADDING_1, CIPHER_1),
+        (PLAIN_PADDING_2, CIPHER_2),
+    ])
+    def test_decrypt(self,plain, cipher):
+        assert self.cipher.decrypt(cipher) == plain

tests/test_decrypt.py

@@ -0,0 +1,62 @@
+from pyencrypt.decrypt import *
+from pyencrypt.encrypt import encrypt_file, encrypt_key
+from constants import AES_KEY
+import pytest
+
+from pyencrypt.generate import generate_aes_key
+
+
+@pytest.mark.parametrize('key', [
+    AES_KEY,
+    generate_aes_key(),
+])
+def test_decrypt_key(key):
+    cipher_key, d, n = encrypt_key(key)
+    assert decrypt_key(cipher_key, d, n) == key.decode()
+
+
+@pytest.fixture
+def encrypted_python_file_path(tmp_path):
+    path = tmp_path / 'test.py'
+    path.touch()
+    path.write_text('print("hello world")')
+    new_path = tmp_path / 'test.pye'
+    encrypt_file(path, AES_KEY, new_path=new_path)
+    path.unlink()
+    return new_path
+
+
+@pytest.mark.parametrize(
+    'path,key,exception', [
+        (Path('tests/test.py'), AES_KEY, Exception),
+        (Path('tests/__init__.pye'), AES_KEY, FileNotFoundError),
+    ]
+)
+def test_decrypt_file_exception(path, key, exception):
+    with pytest.raises(exception) as excinfo:
+        decrypt_file(path, key)
+    assert excinfo.value.__class__ == exception
+
+
+def test_decrypt_file_default(encrypted_python_file_path):
+    assert isinstance(decrypt_file(encrypted_python_file_path, AES_KEY), bytes) == True
+    assert encrypted_python_file_path.exists() == True
+
+
+def test_decrypt_file_delete_origin(encrypted_python_file_path):
+    decrypt_file(encrypted_python_file_path, AES_KEY, delete_origin=True)
+    assert encrypted_python_file_path.exists() == False
+
+
+def test_decrypt_file_new_path(encrypted_python_file_path):
+    new_path = encrypted_python_file_path.parent / 'test.py'
+    decrypt_file(encrypted_python_file_path, AES_KEY, new_path=new_path)
+    assert new_path.exists() == True
+    assert encrypted_python_file_path.exists() == True
+
+
+def test_decrypt_file_new_path_exception(encrypted_python_file_path):
+    new_path = encrypted_python_file_path.parent / 'test.pye'
+    with pytest.raises(Exception) as excinfo:
+        decrypt_file(encrypted_python_file_path, AES_KEY, new_path=new_path)
+    assert str(excinfo.value) == 'Origin file path must be py suffix.'

tests/test_encrypt.py

@@ -0,0 +1,107 @@
+import shutil
+import pytest
+from pyencrypt.encrypt import *
+from pyencrypt.decrypt import *
+import os
+
+from constants import AES_KEY
+from pyencrypt.generate import generate_aes_key
+
+
+
+@pytest.mark.parametrize('key', [
+    AES_KEY,
+    generate_aes_key(),
+])
+def test_encrypt_key(key):
+    cipher, d, n = encrypt_key(key)
+    assert isinstance(cipher, str)
+    assert isinstance(d, int)
+    assert isinstance(n, int)
+
+
+@pytest.mark.parametrize(
+    'path,expected', [
+        (Path('__init__.py'), False),
+        (Path('pyencrypt/__init__.py'), False),
+        (Path('management/commands/user.py'), False),
+        (Path('tests/test.pye'), False),
+        (Path('tests/test_encrypt.py'), True),
+    ]
+)
+def test_can_encrypt(path, expected):
+    assert can_encrypt(path) == expected
+
+
+class TestGenarateSoFile:
+
+    def setup_method(self, method):
+        if method.__name__ == 'test_generate_so_file_default_path':
+            shutil.rmtree((Path(os.getcwd()) / 'encrypted').as_posix(), ignore_errors=True)
+
+    @pytest.mark.parametrize('key', [
+        AES_KEY,
+        generate_aes_key(),
+    ])
+    def test_generate_so_file(self, key, tmp_path):
+        cipher_key, d, n = encrypt_key(key)
+        assert generate_so_file(cipher_key, d, n, tmp_path)
+        assert (tmp_path / 'encrypted' / 'loader.py').exists() == True
+        assert (tmp_path / 'encrypted' / 'loader_origin.py').exists() == True
+        assert list((tmp_path / 'encrypted').glob('loader.cpython-*-*.so')) != []
+
+    @pytest.mark.parametrize('key', [
+        AES_KEY,
+        generate_aes_key(),
+    ])
+    def test_generate_so_file_default_path(self, key):
+        cipher_key, d, n = encrypt_key(key)
+        assert generate_so_file(cipher_key, d, n)
+        assert (Path(os.getcwd()) / 'encrypted' / 'loader.py').exists() == True
+        assert (Path(os.getcwd()) / 'encrypted' / 'loader_origin.py').exists() == True
+        assert list((Path(os.getcwd()) / 'encrypted').glob('loader.cpython-*-*.so')) != []
+
+
+@pytest.mark.parametrize(
+    'path,key,exception',
+    [
+        (Path('tests/test.py'), AES_KEY, FileNotFoundError),
+        (Path('tests/test.pye'), AES_KEY, Exception),  # TODO: 封装Exception
+        (Path('tests/__init__.py'), AES_KEY, Exception),
+    ]
+)
+def test_encrypt_file_path_exception(path, key, exception):
+    with pytest.raises(exception) as excinfo:
+        encrypt_file(path, key)
+    assert excinfo.value.__class__ == exception
+
+
+@pytest.fixture
+def python_file_path(tmp_path):
+    fn = tmp_path / 'test.py'
+    fn.touch()
+    fn.write_text('print("hello world")')
+    return fn
+
+
+def test_encrypt_file_default(python_file_path):
+    assert isinstance(encrypt_file(python_file_path, AES_KEY), bytes) == True
+    assert python_file_path.exists() == True
+
+
+def test_encrypt_file_delete_origin(python_file_path):
+    encrypt_file(python_file_path, AES_KEY, delete_origin=True)
+    assert python_file_path.exists() == False
+
+
+def test_encrypt_file_new_path(python_file_path):
+    new_path = python_file_path.parent / 'test.pye'
+    encrypt_file(python_file_path, AES_KEY, new_path=new_path)
+    assert new_path.exists() == True
+    assert python_file_path.exists() == True
+
+def test_encrypt_file_new_path_exception(python_file_path):
+    new_path = python_file_path.parent / 'test.py'
+    with pytest.raises(Exception) as excinfo:
+        encrypt_file(python_file_path, AES_KEY, new_path=new_path)
+    assert str(excinfo.value) == 'Encrypted file path must be pye suffix.'

tests/test_generate.py

@@ -0,0 +1,31 @@
+from random import randint
+from pyencrypt.generate import *
+import pytest
+
+
+def test_generate_aes_key_default():
+    assert isinstance(generate_aes_key(), bytes)
+
+
+@pytest.mark.parametrize('size', [32, 64, 1024, 4096])
+def test_generate_aes_key(size):
+    assert isinstance(generate_aes_key(size), bytes)
+
+
+@pytest.mark.parametrize('bits', [1024, 2048, 4096])
+def test_generate_rsa_number(bits):
+    numbers = generate_rsa_number(bits)
+    assert len(numbers) == 5
+    p, q, n, e, d = numbers['p'], numbers['q'], numbers['n'], numbers['e'], numbers['d']
+    assert p * q == n
+    assert e * d % (p - 1) == 1
+    assert e * d % (q - 1) == 1
+    plain = randint(0, n)
+    assert pow(pow(plain, e, n), d, n) == plain
+
+
+@pytest.mark.parametrize('bits', [123, 456, 789])
+def test_generate_rsa_number_exception(bits):
+    with pytest.raises(ValueError) as excinfo:
+        generate_rsa_number(bits)
+    assert str(excinfo.value) == "RSA modulus length must be a multiple of 256 and >= 1024"

tests/test_loader.py

@@ -0,0 +1,84 @@
+import pytest
+from pathlib import Path
+import shutil
+from pyencrypt.encrypt import *
+from constants import AES_KEY
+
+
+@pytest.fixture(scope='module')
+def python_file(tmp_path_factory):
+    tmp_path = tmp_path_factory.mktemp('file')
+    path = tmp_path / 'aaa.py'
+    path.touch()
+    path.write_text("""\
+def main():
+    return 'hello world'
+    """)
+    new_path = path.with_suffix('.pye')
+    encrypt_file(path, AES_KEY, new_path=new_path)
+    path.unlink()
+    cipher_key, d, n = encrypt_key(AES_KEY)
+    generate_so_file(cipher_key, d, n, tmp_path)
+    loader_path = list((tmp_path / 'encrypted').glob('loader.cpython-*-*.so'))[0]
+    shutil.copy(loader_path, tmp_path)
+    shutil.rmtree(tmp_path / 'encrypted')
+    return new_path
+
+
+def test_python_file_with_sys_path(python_file: Path, monkeypatch):
+    monkeypatch.syspath_prepend(python_file.parent.as_posix())
+    import loader
+    from aaa import main
+    assert main() == 'hello world'
+
+
+@pytest.fixture(scope='module')
+def python_package(tmp_path_factory):
+    pkg_path = tmp_path_factory.mktemp('package')
+    path = pkg_path / 'bbb'
+    path.mkdir()
+    (path / '__init__.py').touch()
+    path /= 'ccc.py'
+    path.touch()
+    path.write_text("""\
+def main():
+    return 'hello world'
+    """)
+    new_path = path.with_suffix('.pye')
+    encrypt_file(path, AES_KEY, new_path=new_path)
+    path.unlink()
+    cipher_key, d, n = encrypt_key(AES_KEY)
+    generate_so_file(cipher_key, d, n, pkg_path)
+    loader_path = list((pkg_path / 'encrypted').glob('loader.cpython-*-*.so'))[0]
+    shutil.copy(loader_path, pkg_path)
+    shutil.rmtree(pkg_path / 'encrypted')
+    return pkg_path
+
+
+def test_python_package(python_package: Path, monkeypatch):
+    monkeypatch.syspath_prepend(python_package.as_posix())
+    import loader
+    from bbb.ccc import main
+    assert main() == 'hello world'
+
+
+def test_python_package_without_init_file(tmp_path_factory):
+    pkg_path = tmp_path_factory.mktemp('package')
+    path = pkg_path / 'aaa'
+    path.mkdir()
+    path /= 'bbb.py'
+    path.touch()
+    path.write_text("""\
+def main():
+    return 'hello world'
+    """)
+    new_path = path.with_suffix('.pye')
+    encrypt_file(path, AES_KEY, new_path=new_path)
+    path.unlink()
+    cipher_key, d, n = encrypt_key(AES_KEY)
+    generate_so_file(cipher_key, d, n, pkg_path)
+    loader_path = list((pkg_path / 'encrypted').glob('loader.cpython-*-*.so'))[0]
+    shutil.copy(loader_path, pkg_path)
+    shutil.rmtree(pkg_path / 'encrypted')
+    import loader
+    from bbb.ccc import main