Python封装TinyAES

Python封装TinyAES

背景

  偶然发现TinyAES这个项目https://github.com/kokke/tiny-AES-c,本人用Python顺手,就想办法用Python来测试之,于是有了本篇

开发

  参照之前的文章准备好pybind

#include <iostream>
#include <pybind11/pybind11.h>
namespace py = pybind11;

extern "C" {
    #include "aes.h"
    #include "rsa.h"
}

struct TinyAES {
    AES_ctx ctx;
    TinyAES(const std::string& key) {
        AES_init_ctx(&ctx, (uint8_t*)key.data());
    }
    TinyAES(const std::string& key, const std::string& iv) {
        AES_init_ctx_iv(&ctx, (uint8_t*)key.data(), (uint8_t*)iv.data());
    }
    void set_iv(const std::string& iv) {
        AES_ctx_set_iv(&ctx, (uint8_t*)iv.data());
    }
    py::bytes ecb_encrypt(const std::string& buf) {
        std::string tmp = buf; // in-place return
        AES_ECB_encrypt(&ctx, (uint8_t*)tmp.data());
        return py::bytes(tmp.data(), tmp.size());
    }
    py::bytes ecb_decrypt(const std::string& buf) {
        std::string tmp = buf; // in-place return
        AES_ECB_decrypt(&ctx, (uint8_t*)tmp.data());
        return py::bytes(tmp.data(), tmp.size());
    }
    py::bytes cbc_encrypt(const std::string& buf) {
        std::string tmp = buf; // in-place return
        AES_CBC_encrypt_buffer(&ctx, (uint8_t*)tmp.data(), tmp.length());
        return py::bytes(tmp.data(), tmp.size());
    }
    py::bytes cbc_decrypt(const std::string& buf) {
        std::string tmp = buf; // in-place return
        AES_CBC_decrypt_buffer(&ctx, (uint8_t*)tmp.data(), tmp.length());
        return py::bytes(tmp.data(), tmp.size());
    }
};

PYBIND11_MODULE(pyutil, m) {
    py::class_<TinyAES>(m, "TinyAES")
        .def(py::init<const std::string&>())
        .def(py::init<const std::string&,const std::string&>())
        .def("set_iv", &TinyAES::set_iv)
        .def("ecb_encrypt", &TinyAES::ecb_encrypt)
        .def("ecb_decrypt", &TinyAES::ecb_decrypt)
        .def("cbc_encrypt", &TinyAES::cbc_encrypt)
        .def("cbc_decrypt", &TinyAES::cbc_decrypt);
}

测试

import pyutils

key = b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"
iv = b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"
data = b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"
print("origin:", data.hex())
aes = pyutils.TinyAES(key, iv)
data1 = aes.cbc_encrypt(data)
print("encrypt:", data1.hex())
aes = pyutils.TinyAES(key, iv)
data2 = aes.cbc_decrypt(data1)
print("decrypt:", data2.hex())

得到:

origin: 00000000000000000000000000000000
encrypt: 66e94bd4ef8a2c3b884cfa59ca342b2e
decrypt: 00000000000000000000000000000000