diff --git a/pyencrypt/encrypt.py b/pyencrypt/encrypt.py index af3a0b7..1c328d3 100644 --- a/pyencrypt/encrypt.py +++ b/pyencrypt/encrypt.py @@ -1,4 +1,5 @@ import os +import sys import re from pathlib import Path from typing import Optional @@ -97,8 +98,18 @@ def generate_so_file(cipher_key: str, d: int, n: int, base_dir: Optional[Path] = script_args=['build_ext', '--build-lib', temp_dir.as_posix()], cmdclass={'build_ext': build_ext}, ) - print(list(temp_dir.iterdir())) - return list(temp_dir.glob('loader.cpython-*-*.so'))[0].absolute() + if sys.platform.startswith('win'): + # loader.cp36-win_amd64.pyd + pattern = 'loader.cp*-*.pyd' + else: + # loader.cpython-36m-x86_64-linux-gnu.so + # loader.cpython-36m-darwin.so + pattern = "loader.cpython-*-*.so" + + loader_extension = next(temp_dir.glob(pattern), None) + if loader_extension is None: + raise Exception(f"Can't find loader extension in {temp_dir.as_posix()}") + return loader_extension.absolute() def encrypt_file(path: Path, key: str, delete_origin: bool = False, new_path: Optional[Path] = None): diff --git a/tests/test_encrypt.py b/tests/test_encrypt.py index f5f7736..c0e51cc 100644 --- a/tests/test_encrypt.py +++ b/tests/test_encrypt.py @@ -1,6 +1,7 @@ import os from pathlib import Path import shutil +import sys import pytest from pyencrypt.encrypt import can_encrypt, encrypt_file, encrypt_key, generate_so_file @@ -58,7 +59,10 @@ def test_generate_so_file_default_path(self, key): assert generate_so_file(cipher_key, d, n) assert (Path(os.getcwd()) / 'encrypted' / 'loader.py').exists() is True assert (Path(os.getcwd()) / 'encrypted' / 'loader_origin.py').exists() is True - assert list((Path(os.getcwd()) / 'encrypted').glob('loader.cpython-*-*.so')) != [] + if sys.platform.startswith('win'): + assert next((Path(os.getcwd()) / 'encrypted').glob('loader.cp*-*.pyd'), None) is not None + else: + assert next((Path(os.getcwd()) / 'encrypted').glob('loader.cpython-*-*.so'), None) is not None @pytest.mark.parametrize(