Update unittest.yml

.github/workflows/unittest.yml

@@ -32,7 +32,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.11", "3.12"]
+        python-version: ["3.12"]
+        gmssl-version: ["v3.1.1", "v3.1.0"]
         os: [ubuntu-latest, macos-latest, windows-latest]
 
     steps:
@@ -43,7 +44,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Download GmSSL
       run: |
-        git clone -b 'v3.1.1' --depth 1 https://github.com/guanzhi/GmSSL.git
+        git clone -b '${{ matrix.gmssl-version }}' --depth 1 https://github.com/guanzhi/GmSSL.git
     - name: Install GmSSL on ubuntu and macos
       if: ${{ startsWith(matrix.os, 'ubuntu') || startsWith(matrix.os, 'macos') }}
       working-directory: ${{ github.workspace }}/GmSSL
@@ -71,7 +72,7 @@ jobs:
     - name: Install Python dependencies
       run: |
         python -m pip install poetry
-        poetry install --sync --with dev
+        poetry install --sync --without dev
     - name: Test with pytest
       run: |
         poetry run pytest --cov=src --cov-report=xml

pyproject.toml

@@ -26,6 +26,8 @@ pycryptodomex = "^3.20.0"
 autopep8 = "^2.1.0"
 isort = "^5.13.2"
 flake8 = "^7.0.0"
+
+[tool.poetry.group.test.dependencies]
 pytest = "^8.1.1"
 pytest-cov = "^5.0.0"
 

src/pygmssl/_gm.py

@@ -1,12 +1,20 @@
 import sys
 import warnings
-from ctypes import cdll, c_char_p
+from ctypes import cdll, c_char_p, c_void_p, c_int
 from ctypes.util import find_library
 
 if sys.platform == 'win32':  # pragma: no cover
     libc = cdll.LoadLibrary(find_library('msvcrt'))
+    win32 = True
 else:
     libc = cdll.LoadLibrary(find_library('c'))
+    win32 = False
+
+libc.fopen.argtypes = [c_char_p, c_char_p]
+libc.fopen.restype = c_void_p
+libc.fclose.argtypes = [c_void_p]
+libc.fclose.restype = c_int
+
 
 libgm = find_library('gmssl')
 
@@ -17,3 +25,11 @@
 else:
     _gm = cdll.LoadLibrary(libgm)
     _gm.gmssl_version_str.restype = c_char_p
+    _gm.sm2_private_key_info_encrypt_to_pem.argtypes = [c_void_p, c_char_p, c_void_p]
+    _gm.sm2_private_key_info_encrypt_to_pem.restype = c_int
+    _gm.sm2_private_key_info_decrypt_from_pem.argtypes = [c_void_p, c_char_p, c_void_p]
+    _gm.sm2_private_key_info_decrypt_from_pem.restype = c_int
+    _gm.sm2_public_key_info_from_pem.argtypes = [c_void_p, c_void_p]
+    _gm.sm2_public_key_info_from_pem.restype = c_int
+    _gm.sm2_public_key_info_to_pem.argtypes = [c_void_p, c_void_p]
+    _gm.sm2_public_key_info_to_pem.restype = c_int

src/pygmssl/sm2.py

@@ -1,9 +1,11 @@
-from ctypes import byref, c_uint8, c_size_t, Structure, c_char_p, c_void_p
+import base64
+from ctypes import byref, c_uint8, c_size_t, Structure, c_char_p, pointer
 import tempfile
+import os
 
 from Cryptodome.Util.asn1 import DerSequence
 
-from ._gm import _gm, libc
+from ._gm import _gm, libc, win32
 from .sm3 import _SM3CTX
 
 SM2_DEFAULT_ID = b'1234567812345678'
@@ -129,46 +131,127 @@ def decrypt(self, data: bytes) -> bytes:
         _gm.sm2_decrypt(byref(self._sm2_key), byref(buff), len(data), byref(out), byref(length))
         return bytes(out[:length.value])
 
-    def export_encrypted_private_key_to_pem(self, password: bytes) -> bytes:
-        with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
-            libc.fopen.restype = c_void_p
-            fp = libc.fopen(tmp_f.name.encode('utf8'), 'wb')
-            assert _gm.sm2_private_key_info_encrypt_to_pem(byref(self._sm2_key), c_char_p(password), c_void_p(fp)) == 1
-            libc.fclose(c_void_p(fp))
-            with open(tmp_f.name, 'rb') as f:
+    def _export_encrypted_pri_to_der(self, password: bytes) -> bytes:
+        buff = (c_uint8 * 4096)()
+        length = c_size_t()
+        _gm.sm2_private_key_info_encrypt_to_der(byref(self._sm2_key), password, byref(pointer(buff)), byref(length))
+        return bytes(buff[:length.value])
+
+    def _export_pub_to_der(self) -> bytes:
+        buff = (c_uint8 * 4096)()
+        length = c_size_t()
+        _gm.sm2_public_key_info_to_der(byref(self._sm2_key), byref(pointer(buff)), byref(length))
+        return bytes(buff[:length.value])
+
+    def _nix_export_private_key_to_encrypted_pem(self, password: bytes) -> bytes:
+        with tempfile.NamedTemporaryFile(delete=False) as _tmp_f:
+            tmp_f_name = _tmp_f.name
+            fp = libc.fopen(tmp_f_name.encode('utf8'), b'wb')
+
+            assert _gm.sm2_private_key_info_encrypt_to_pem(byref(self._sm2_key), c_char_p(password), fp) == 1
+            libc.fclose(fp)
+            with open(tmp_f_name, 'rb') as f:
                 res = f.read()
             return res
 
-    def export_public_key_to_pem(self) -> bytes:
-        with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
-            libc.fopen.restype = c_void_p
-            fp = libc.fopen(tmp_f.name.encode('utf8'), 'wb')
-            assert _gm.sm2_public_key_info_to_pem(byref(self._sm2_key), c_void_p(fp)) == 1
-            libc.fclose(c_void_p(fp))
-            with open(tmp_f.name, 'rb') as f:
+    def _win_export_private_key_to_encrypted_pem(self, password: bytes) -> bytes:
+        der = self._export_encrypted_pri_to_der(password)
+        return self._pem_write(der, 'ENCRYPTED PRIVATE KEY')
+
+    def _nix_export_public_key_to_pem(self) -> bytes:
+        with tempfile.NamedTemporaryFile(delete=False) as _tmp_f:
+            tmp_f_name = _tmp_f.name
+            fp = libc.fopen(tmp_f_name.encode('utf8'), b'wb')
+            assert _gm.sm2_public_key_info_to_pem(byref(self._sm2_key), fp) == 1
+            libc.fclose(fp)
+            with open(tmp_f_name, 'rb') as f:
                 res = f.read()
             return res
 
+    def _win_export_public_key_to_pem(self) -> bytes:
+        pub_der = self._export_pub_to_der()
+        return self._pem_write(pub_der, 'PUBLIC KEY')
+
+    def _pem_write(self, der: bytes, name: str) -> bytes:
+        data = base64.b64encode(der).decode('utf8')
+        prefix = f'-----BEGIN {name}-----'
+        suffix = f'-----END {name}-----'
+        tmp: list[str] = [prefix]
+        for i in range(0, len(data), 64):
+            chunk = data[i:i + 64]
+            tmp.append(chunk)
+        tmp.append(suffix)
+        return ''.join(_line + os.linesep for _line in tmp).encode('utf8')
+
     @classmethod
-    def import_private_from_pem(cls, pem: bytes, password: bytes) -> 'SM2':
-        with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
-            with open(tmp_f.name, 'wb') as f:
+    def _nix_import_private_key_from_encrypted_pem(cls, pem: bytes, password: bytes) -> 'SM2':
+        with tempfile.NamedTemporaryFile(delete=False) as _tmp_f:
+            tmp_f_name = _tmp_f.name
+            with open(tmp_f_name, 'wb') as f:
                 f.write(pem)
-            libc.fopen.restype = c_void_p
-            fp = libc.fopen(tmp_f.name.encode('utf8'), 'rb')
+            fp = libc.fopen(tmp_f_name.encode('utf8'), b'rb')
             obj = SM2()
-            assert _gm.sm2_private_key_info_decrypt_from_pem(byref(obj._sm2_key), c_char_p(password), c_void_p(fp)) == 1
-            libc.fclose(c_void_p(fp))
+            assert _gm.sm2_private_key_info_decrypt_from_pem(byref(obj._sm2_key), c_char_p(password), fp) == 1
+            libc.fclose(fp)
             return obj
 
     @classmethod
-    def import_public_from_pem(cls, pem: bytes) -> 'SM2':
-        with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
-            with open(tmp_f.name, 'wb') as f:
+    def _nix_import_public_key_from_pem(cls, pem: bytes) -> 'SM2':
+        with tempfile.NamedTemporaryFile(delete=False) as _tmp_f:
+            tmp_f_name = _tmp_f.name
+            with open(tmp_f_name, 'wb') as f:
                 f.write(pem)
-            libc.fopen.restype = c_void_p
-            fp = libc.fopen(tmp_f.name.encode('utf8'), 'rb')
+            fp = libc.fopen(tmp_f_name.encode('utf8'), b'rb')
             obj = SM2()
-            assert _gm.sm2_public_key_info_from_pem(byref(obj._sm2_key), c_void_p(fp)) == 1
-            libc.fclose(c_void_p(fp))
+            assert _gm.sm2_public_key_info_from_pem(byref(obj._sm2_key), fp) == 1
+            libc.fclose(fp)
             return obj
+
+    @staticmethod
+    def _pem_read(pem: str, name: str) -> bytes:
+        tmp = pem.splitlines()
+        prefix = f'-----BEGIN {name}-----'
+        suffix = f'-----END {name}-----'
+        assert tmp[0] == prefix
+        assert tmp[-1] == suffix
+        mid = ''.join(tmp[1:-1])
+        return base64.b64decode((mid + ('=' * (-len(mid) % 4))).encode())
+
+    @classmethod
+    def _win_import_private_key_from_encrypted_pem(cls, pem: bytes, password: bytes) -> 'SM2':
+        der_data = cls._pem_read(pem.decode('utf8'), 'ENCRYPTED PRIVATE KEY')
+        obj = SM2()
+        attr = (c_uint8 * 4096)()
+        attr_len = c_size_t()
+        p = pointer(attr)
+        buf = (c_uint8 * 4096)()
+        buf[:len(der_data)] = der_data
+        buflen = c_size_t(len(der_data))
+        cp = pointer(buf)
+        assert _gm.sm2_private_key_info_decrypt_from_der(byref(obj._sm2_key), byref(
+            p), byref(attr_len), password, byref(cp), byref(buflen)) == 1
+        assert buflen.value == 0
+        return obj
+
+    @classmethod
+    def _win_import_public_key_from_pem(cls, pem: bytes) -> 'SM2':
+        der_data = cls._pem_read(pem.decode('utf8'), 'PUBLIC KEY')
+        obj = SM2()
+        buf = (c_uint8 * 4096)()
+        buf[:len(der_data)] = der_data
+        vlen = c_size_t(len(der_data))
+        cp = pointer(buf)
+        assert _gm.sm2_public_key_info_from_der(byref(obj._sm2_key), byref(cp), byref(vlen)) == 1
+        assert vlen.value == 0
+        return obj
+
+    if win32:
+        export_public_key_to_pem = _win_export_public_key_to_pem
+        export_private_key_to_encrypted_pem = _win_export_private_key_to_encrypted_pem
+        import_public_key_from_pem = _win_import_public_key_from_pem
+        import_private_key_from_encrypted_pem = _win_import_private_key_from_encrypted_pem
+    else:
+        export_public_key_to_pem = _nix_export_public_key_to_pem
+        export_private_key_to_encrypted_pem = _nix_export_private_key_to_encrypted_pem
+        import_public_key_from_pem = _nix_import_public_key_from_pem
+        import_private_key_from_encrypted_pem = _nix_import_private_key_from_encrypted_pem

tests/test_sm2.py

@@ -73,31 +73,34 @@ def test_010_sm2_sign_with_id_asn1(self):
         self.assertFalse(self.k.verify(data + b'\x00', sig, id=b'test', asn1=True))
         self.assertTrue(self.k.verify(data, sig, id=b'test', asn1=True))
 
-    def test_private_pem_export_and_import(self):
+    def test_100_private_pem_export_and_import(self):
         password = b'test-123-456'
         obj = SM2.generate_new_pair()
         assert obj.pub_key != b'\x00' * 64
         assert obj.pri_key != b'\x00' * 32
-        new_obj = SM2.import_private_from_pem(obj.export_encrypted_private_key_to_pem(password), password)
+        pem = obj.export_private_key_to_encrypted_pem(password)
+        new_obj = SM2.import_private_key_from_encrypted_pem(pem, password)
         assert new_obj.pri_key != b'\x00' * 32
         assert new_obj.pri_key == obj.pri_key
 
         assert new_obj.pub_key != b'\x00' * 64
         assert new_obj.pub_key == obj.pub_key
 
-    def test_pub_pem_export_and_import(self):
+    def test_101_pub_pem_export_and_import(self):
         obj = SM2.generate_new_pair()
         assert obj.pub_key != b'\x00' * 64
         assert obj.pri_key != b'\x00' * 32
-        new_obj = SM2.import_public_from_pem(obj.export_public_key_to_pem())
+        pem = obj.export_public_key_to_pem()
+        new_obj = SM2.import_public_key_from_pem(pem)
         assert new_obj.pri_key == b'\x00' * 32
         assert new_obj.pub_key != b'\x00' * 64
         assert new_obj.pub_key == obj.pub_key
 
-    def test_error_import_private_pem(self):
+    def test_102_error_import_private_pem(self):
         password = b'test-123-456'
         obj = SM2.generate_new_pair()
         assert obj.pub_key != b'\x00' * 64
         assert obj.pri_key != b'\x00' * 32
         with self.assertRaises(Exception):
-            SM2.import_private_from_pem(obj.export_encrypted_private_key_to_pem(password), b'wrong-password')
+            pem = obj.export_private_key_to_encrypted_pem(password)
+            SM2.import_private_key_from_encrypted_pem(pem, b'wrong-password')