module.cpp

#define PY_SSIZE_T_CLEAN
#include <Python.h>

#include "config.h"
#include "base64.hpp"
#include "gzip.hpp"

#include <iostream>
#include <cstdint>
#include <string>
#include <gmlib/rng/std_rng.h>
#include <gmlib/sm2/sm2.h>
#include <gmlib/sm3/sm3.h>
#include <gmlib/sm4/sm4.h>
#include <gmlib/sm4/sm4_mode.h>

using namespace Pectics;

constexpr size_t BUFFER_SIZE = 1024 * 1024; // 1MB

// P3s0+rMuY4Nt5cUWuOCjMhDzVNdom+W0RvdV6ngM+/E=
constexpr uint8_t PRIVATE_KEY[32] = {
	0x3f, 0x7b, 0x34, 0xfa, 0xb3, 0x2e, 0x63, 0x83,
	0x6d, 0xe5, 0xc5, 0x16, 0xb8, 0xe0, 0xa3, 0x32,
	0x10, 0xf3, 0x54, 0xd7, 0x68, 0x9b, 0xe5, 0xb4,
	0x46, 0xf7, 0x55, 0xea, 0x78, 0x0c, 0xfb, 0xf1,
};

// BL7JvEAV7Wci0h5YAysN0BPNVdcUhuyJszJLRwnurav0CGftcrVcvrWeCPBIjIIBF371teRbrCS9V1Wyq7i3Arc=
constexpr uint8_t PUBLIC_KEY[2][32] = { {
	0xbe, 0xc9, 0xbc, 0x40, 0x15, 0xed, 0x67, 0x22,
	0xd2, 0x1e, 0x58, 0x03, 0x2b, 0x0d, 0xd0, 0x13,
	0xcd, 0x55, 0xd7, 0x14, 0x86, 0xec, 0x89, 0xb3,
	0x32, 0x4b, 0x47, 0x09, 0xee, 0xad, 0xab, 0xf4,
}, {
	0x08, 0x67, 0xed, 0x72, 0xb5, 0x5c, 0xbe, 0xb5,
	0x9e, 0x08, 0xf0, 0x48, 0x8c, 0x82, 0x01, 0x17,
	0x7e, 0xf5, 0xb5, 0xe4, 0x5b, 0xac, 0x24, 0xbd,
	0x57, 0x55, 0xb2, 0xab, 0xb8, 0xb7, 0x02, 0xb7,
} };

static PyObject* C_SM2Encrypt(PyObject*, PyObject* o) {

	// Check input
	if (!PyUnicode_Check(o)) {
		PyErr_SetString(PyExc_UnicodeError, "Input is not a string");
		return _Py_NULL;
	}

	// Parse plain text (Plain)
	std::string text = PyBytes_AsString(PyUnicode_AsUTF8String(o));
	size_t plain_len = text.size();
	uint8_t* plain = new uint8_t[plain_len];
	std::copy(text.begin(), text.end(), plain);

	// Initialize public key
	sm2::SM2PublicKey<sm3::SM3> pub_key(PUBLIC_KEY[0], PUBLIC_KEY[1]);

	// Check buffer size
	if (pub_key.ciphertext_len(plain, plain_len) > BUFFER_SIZE) {
		PyErr_SetString(PyExc_BufferError, "Buffer overflow");
		return _Py_NULL;
	}

	// Encrypt cipher (Binary)
	uint8_t* cipher = new uint8_t[BUFFER_SIZE];
	size_t cipher_len;
	rng::StdRng rng;
	pub_key.encrypt(cipher, &cipher_len, plain, plain_len, rng);

	// Encode result (Base64)
	size_t result_len = (cipher_len + 2) / 3 * 4;
	uint8_t* result = new uint8_t[result_len];
	Base64::Encode(cipher, cipher_len, result);

	// Return
	std::string ret(result, result + result_len);
	delete[] plain;
	delete[] cipher;
	delete[] result;
	return PyUnicode_FromStringAndSize(ret.data(), ret.size());
}

static PyObject* C_SM2Decrypt(PyObject*, PyObject* o) {

	// Check input
	if (!PyUnicode_Check(o)) {
		PyErr_SetString(PyExc_UnicodeError, "Input is not a string");
		return _Py_NULL;
	}

	// Parse cipher text (Base64)
	std::string text = PyBytes_AsString(PyUnicode_AsUTF8String(o));
	size_t cipher_len = text.size() / 4 * 3;
	if (text.size() >= 1 && text[text.size() - 1] == '=') --cipher_len;
	if (text.size() >= 2 && text[text.size() - 2] == '=') --cipher_len;

	// Decode cipher (Binary)
	uint8_t* cipher = new uint8_t[cipher_len];
	Base64::Decode(reinterpret_cast<const uint8_t*>(text.data()), text.size(), cipher);

	// Initialize private key
	sm2::SM2PrivateKey<sm3::SM3> pri_key(PRIVATE_KEY);

	// Check buffer size
	if (pri_key.plaintext_len(cipher, cipher_len) > BUFFER_SIZE) {
		PyErr_SetString(PyExc_BufferError, "Buffer overflow");
		return _Py_NULL;
	}

	// Decrypt plain (Plain)
	uint8_t* plain = new uint8_t[BUFFER_SIZE];
	size_t plain_len;
	pri_key.decrypt(plain, &plain_len, cipher, cipher_len);

	// Return
	std::string ret(plain, plain + plain_len);
	delete[] plain;
	delete[] cipher;
	return PyUnicode_FromStringAndSize(ret.data(), ret.size());
}

static PyObject* C_SM4Encrypt(PyObject*, PyObject* o) {

	// Check input
	if (!PyTuple_Check(o))
		return _Py_NULL;

	// Parse keywords
	const char* t;
	const char* k;
	bool gzip = false;
	if (!PyArg_ParseTuple(o, "ss|i", &t, &k, &gzip))
		return _Py_NULL;

	// Parse text (Plain)
	uint8_t* text = reinterpret_cast<uint8_t*>(const_cast<char*>(t));
	size_t len = strlen(t);
	uint8_t* key_str = reinterpret_cast<uint8_t*>(const_cast<char*>(k));
	size_t key_len = strlen(k);

	// Parse key
	uint8_t key[sm4::SM4::USER_KEY_LEN];
	Base64::Decode(key_str, key_len, key);

	// Initialize encryptor
	sm4::SM4EcbEncryptor enc(key);

	// Compress
	if (gzip) GZip::Compress(text, len, text, len);

	// Parse plain (Plain/Compressed) (Padded)
	size_t block_size = sm4::SM4::BLOCK_SIZE;
	size_t padded_len = block_size * (len / block_size + 1);
	uint8_t* plain = new uint8_t[padded_len];
	std::copy(text, text + len, plain);
	std::fill(plain + len, plain + padded_len, padded_len - len);

	// Encrypt cipher (Binary)
	uint8_t* cipher = new uint8_t[BUFFER_SIZE];
	size_t cipher_len;
	enc.do_final(cipher, &cipher_len, plain, padded_len);

	// Encode result (Base64)
	size_t result_len = (cipher_len + 2) / 3 * 4;
	uint8_t* result = new uint8_t[result_len];
	Base64::Encode(cipher, cipher_len, result);

	// Return
	std::string ret(result, result + result_len);
	delete[] plain;
	delete[] cipher;
	return PyUnicode_FromStringAndSize(ret.data(), ret.size());
}

static PyObject* C_SM4Decrypt(PyObject*, PyObject* o) {

	// Check input
	if (!PyTuple_Check(o))
		return _Py_NULL;

	// Parse keywords
	const char* t;
	const char* k;
	bool gzip = false;
	if (!PyArg_ParseTuple(o, "ss|i", &t, &k, &gzip))
		return _Py_NULL;

	// Parse text (Base64)
	uint8_t* text = reinterpret_cast<uint8_t*>(const_cast<char*>(t));
	size_t len = strlen(t);
	uint8_t* key_str = reinterpret_cast<uint8_t*>(const_cast<char*>(k));
	size_t key_len = strlen(k);

	// Parse key
	uint8_t key[sm4::SM4::USER_KEY_LEN];
	Base64::Decode(key_str, key_len, key);

	// Initialize decryptor
	sm4::SM4EcbDecryptor dec(key);

	// Parse cipher (Binary)
	size_t cipher_len = len / 4 * 3;
	if (len >= 1 && text[len - 1] == '=') --cipher_len;
	if (len >= 2 && text[len - 2] == '=') --cipher_len;
	uint8_t* cipher = new uint8_t[cipher_len];
	Base64::Decode(text, len, cipher);

	// Decrypt plain (Plain/Compressed) (Unpadded)
	uint8_t* plain = new uint8_t[BUFFER_SIZE];
	size_t plain_len;
	dec.do_final(plain, &plain_len, cipher, cipher_len);
	plain_len -= plain[plain_len - 1];

	// Decompress
	if (gzip) GZip::Decompress(plain, plain_len, plain, plain_len);

	// Return
	std::string ret(plain, plain + plain_len);
	delete[] plain;
	delete[] cipher;
	return PyUnicode_FromStringAndSize(ret.data(), ret.size());
}

static PyMethodDef methods[] = {
	{ "SM2Encrypt", reinterpret_cast<PyCFunction>(C_SM2Encrypt), METH_O, "Encrypt text with SM2" },
	{ "SM2Decrypt", reinterpret_cast<PyCFunction>(C_SM2Decrypt), METH_O, "Decrypt text with SM2" },
	{ "SM4Encrypt", reinterpret_cast<PyCFunction>(C_SM4Encrypt), METH_VARARGS, "Encrypt text with SM4" },
	{ "SM4Decrypt", reinterpret_cast<PyCFunction>(C_SM4Decrypt), METH_VARARGS, "Decrypt text with SM4" },
	{ nullptr, nullptr, 0, nullptr },
};

static PyModuleDef module = {
	PyModuleDef_HEAD_INIT,
	"CryptUtils",
	"SM2/SM4 encryption and decryption for YUN",
	0,
	methods,
};

PyMODINIT_FUNC PyInit_CryptUtils(void) {
	return PyModule_Create(&module);
}