【PQC】PQMagic(国内首个高性能后量子密码算法库)的Python封装

上一篇文章当中，我们介绍了PQMagic[1]这一个库。然后呢，想着大多数读者，可能会采用Python，所以呢，这里参考liboqs-python[2]的思路，简单的搞一个Python的Wrapper出来，这里简单的记录一下，我个人的设计思路，仅代表个人看法，也欢迎读者和我交流。

前置知识

PQMagic是PQMagic团队推出了国内首个支持FIPS 203、FIPS 204以及FIPS 205标准的高性能后量子密码算法库——PQMagic[1]，该算法库同时还支持PQMagic团队自研的后量子密码算法，并进一步对所有算法使用的哈希模块完成国密改造，实现完全自主可控，提供抗量子威胁能力。PQMagic在兼顾可扩展性和多样化的同时，达到了目前最优的性能表现。

这里，我们采用cookie-cutter[3]来生成，这个框架如何用，这里就不过多赘述了，相信写过Python的package的应该对这个框架并不陌生。

封装设计

这里，我们采用面向对象的设计思路，简单对于这个封装进行一个分层。

基础架构层(采用ctypes封装C接口)

import ctypesimport osclassPQMagicLibrary:def__init__(self, lib_path):        self.lib = ctypes.CDLL(lib_path)        self._register_functions()def_register_functions(self):# 注册C函数，这里以ML-DSA-44的密钥生成为例        self.lib.pqmagic_ml_dsa_44_std_keypair.argtypes = [            ctypes.POINTER(ctypes.c_ubyte),            ctypes.POINTER(ctypes.c_ubyte)        ]        self.lib.pqmagic_ml_dsa_44_std_keypair.restype = ctypes.c_int# 其他函数类似注册...classCryptoBuffer:def__init__(self, length):        self.buffer = (ctypes.c_ubyte * length)()        self.length = lengthdefto_bytes(self):return bytes(self.buffer)    @classmethoddeffrom_bytes(cls, data):        buf = cls(len(data))        ctypes.memmove(buf.buffer, data, len(data))return buf

算法抽象层

这里，我们主要是需要封装两个算法，一个是密钥封装机制，一个是数字签名算法，分别给出具体的接口，方便后面进行扩展，这里我偷懒了，目前签名算法不支持传入ctx, 后面再补充吧。

from abc import ABC, abstractmethodclassPQCryptoAlgorithm(ABC):def__init__(self, lib):        self.lib = lib    @property    @abstractmethoddefCRYPTO_PUBLICKEYBYTES(self):pass    @property    @abstractmethoddefCRYPTO_SECRETKEYBYTES(self):pass    @abstractmethoddefcrypto_keypair(self):passclassDigitalSignatureAlgorithm(PQCryptoAlgorithm):    @property    @abstractmethoddefCRYPTO_SIGBYTES(self):pass    @abstractmethoddefcrypto_sign(self, msg, sk):pass    @abstractmethoddefcrypto_verify(self, sig, msg, pk):passclassKEMAlgorithm(PQCryptoAlgorithm):    @property    @abstractmethoddefCRYPTO_CIPHERTEXTBYTES(self):pass    @property    @abstractmethoddefCRYPTO_SSBYTES(self):pass    @abstractmethoddefcrypto_encrypt(self, pk):pass    @abstractmethoddefcrypto_decrypt(self, ct, sk):pass

算法实现层

针对于需要封装的算法，我们给出具体的实现类，也就是需要实现上面对应的接口，这里，我们还是用ML-DSA-44为例，最后代码会完整给出。

classMLDSA44(DigitalSignatureAlgorithm):    CRYPTO_PUBLICKEYBYTES = 1312    CRYPTO_SECRETKEYBYTES = 2560    CRYPTO_SIGBYTES = 2420defcrypto_keypair(self):        pk = CryptoBuffer(self.CRYPTO_PUBLICKEYBYTES)        sk = CryptoBuffer(self.CRYPTO_SECRETKEYBYTES)        rc = self.lib.pqmagic_ml_dsa_44_std_keypair(            pk.buffer, sk.buffer        )if rc != 0:raise Exception(f"Keypair generation failed with code {rc}")return pk.to_bytes(), sk.to_bytes()defcrypto_sign(self, msg, sk):        sk_buf = CryptoBuffer.from_bytes(sk)        msg_buf = CryptoBuffer.from_bytes(msg)        sig = CryptoBuffer(len(msg) + self.CRYPTO_SIGBYTES)        smlen = ctypes.c_size_t()        rc = self.lib.pqmagic_ml_dsa_44_std(            sig.buffer, ctypes.byref(smlen),            msg_buf.buffer, ctypes.c_ulong(msg_buf.length),None, ctypes.c_ulong(0),            sk_buf.buffer,        )if rc != 0:raise Exception(f"Signing failed with code {rc}")return sig.to_bytes()defcrypto_verify(self, sig, msg, pk):        pk_buf = CryptoBuffer.from_bytes(pk)        sig_buf = CryptoBuffer.from_bytes(sig)        m2 = CryptoBuffer(len(msg) + self.CRYPTO_SIGBYTES)        mlen = ctypes.c_size_t()        rc = self.lib.pqmagic_ml_dsa_44_std_open(            m2.buffer, ctypes.byref(mlen),            sig_buf.buffer, ctypes.c_ulong(sig_buf.length),None, ctypes.c_ulong(0),            pk_buf.buffer        )if rc < 0:raise Exception(f"Verification error with code {rc}")return bool(not rc)

对外接口

这里，参考Java当中的密码学库，我们采用工厂模式，传入具体算法的名称，来自动查找对应的算法。

"""Main module."""from src.pqmagic_python.core.lib import PQMagicLibraryfrom src.pqmagic_python.kem import *from src.pqmagic_python.sig import *classCryptoFactory:    algorithms = {'ML-DSA-44': MLDSA44,'ML-DSA-65': MLDSA65,'ML-DSA-87': MLDSA87,'SLH-DSA-SHA256-128f': SLHDSASHA256_128f,'SLH-DSA-SHA256-128s': SLHDSASHA256_128s,'SLH-DSA-SHA256-192f': SLHDSASHA256_192f,'SLH-DSA-SHA256-192s': SLHDSASHA256_192s,'SLH-DSA-SHA256-256f': SLHDSASHA256_256f,'SLH-DSA-SHA256-256s': SLHDSASHA256_256s,'SLH-DSA-SHAKE-128f': SLHDSASHAKE_128f,'SLH-DSA-SHAKE-128s': SLHDSASHAKE_128s,'SLH-DSA-SHAKE-192f': SLHDSASHAKE_192f,'SLH-DSA-SHAKE-192s': SLHDSASHAKE_192s,'SLH-DSA-SHAKE-256f': SLHDSASHAKE_256f,'SLH-DSA-SHAKE-256s': SLHDSASHAKE_256s,'SLH-DSA-SM3-128f': SLHDSASM3_128f,'SLH-DSA-SM3-128s': SLHDSASM3_128s,'AIGIS-SIG-1': AIGIS_SIG_1,'AIGIS-SIG-2': AIGIS_SIG_2,'AIGIS-SIG-3': AIGIS_SIG_3,'Dilithium2': Dilithium2,'Dilithium3': Dilithium3,'Dilithium5': Dilithium5,'SPHINCS-A-SHA2-128f': SPHINCS_A_SHA2_128f,'SPHINCS-A-SHA2-128s': SPHINCS_A_SHA2_128s,'SPHINCS-A-SHA2-192f': SPHINCS_A_SHA2_192f,'SPHINCS-A-SHA2-192s': SPHINCS_A_SHA2_192s,'SPHINCS-A-SHA2-256f': SPHINCS_A_SHA2_256f,'SPHINCS-A-SHA2-256s': SPHINCS_A_SHA2_256s,'SPHINCS-A-SHAKE-128f': SPHINCS_A_SHAKE_128f,'SPHINCS-A-SHAKE-128s': SPHINCS_A_SHAKE_128s,'SPHINCS-A-SHAKE-192f': SPHINCS_A_SHAKE_192f,'SPHINCS-A-SHAKE-192s': SPHINCS_A_SHAKE_192s,'SPHINCS-A-SHAKE-256f': SPHINCS_A_SHAKE_256f,'SPHINCS-A-SHAKE-256s': SPHINCS_A_SHAKE_256s,'SPHINCS-A-SM3-128f': SPHINCS_A_SM3_128f,'SPHINCS-A-SM3-128s': SPHINCS_A_SM3_128s,'ML-KEM-512': MLKEM512,'ML-KEM-768': MLKEM768,'ML-KEM-1024': MLKEM1024,'Kyber-512': Kyber512,'Kyber-768': Kyber768,'Kyber-1024': Kyber1024,'Aigis-Enc-1': AigisEnc1,'Aigis-Enc-2': AigisEnc2,'Aigis-Enc-3': AigisEnc3,'Aigis-Enc-4': AigisEnc4,    }def__init__(self, lib_path):        self.lib = PQMagicLibrary(lib_path)defcreate(self, algorithm_name):        cls = self.algorithms.get(algorithm_name)ifnot cls:raise ValueError(f"Unsupported algorithm: {algorithm_name}")return cls(self.lib.lib)

目前，PQMagic所支持的库，这里就都适配好了。

正确性测试

这里，我们简单的写一下单元测试，测试封装的正确性就好了。

from src.pqmagic_python.pqmagic_python import CryptoFactoryfactory = CryptoFactory("/usr/local/lib/libpqmagic.so")deftest_ml_dsa_44_correctness():    instance = factory.create('ML-DSA-44')    pk, sk = instance.crypto_keypair()    message = b"Important message"    signature = instance.crypto_sign(message, sk)assert instance.crypto_verify(signature, message, pk)

其他测试代码，可以参考test/test_pqmagic_python.py里面 内容。

可以发现非常的nice，所有测试样例都通过了。

使用样例

这里，简单给出一个使用的样例，我们还是以签名为例。

# 使用示例import osfrom src.pqmagic_python.pqmagic_python import CryptoFactoryif __name__ == "__main__":    factory = CryptoFactory("/usr/local/lib/libpqmagic.so")    instance = factory.create('SLH-DSA-SM3-128s')    pk, sk = instance.crypto_keypair()    message = b"Important message"    signature = instance.crypto_sign(message, sk)assert instance.crypto_verify(signature, message, pk)

其他例子，可以参考examples里的代码。

总结

这里，再次感谢PQMagic团队的开源，最终封装的Python版本，我也放到了github[4]上，有需要的读者可以自行查看源码[4]，本次封装的模式，仅代表我个人的一个思路, 好了，快乐的时光过得特别快，又到了说再见的时候了，咱们下次再见。

参考资料

1. https://pqcrypto.cn/ or https://pqcrypto.dev/
2. https://github.com/open-quantum-safe/liboqs-python
3. https://github.com/audreyr/cookiecutter
4. https://github.com/Litt1eQ/pqmagic_python.git