【Python】Python3调用C++代码之封装TinyAES

lichao

背景

  在分析某款iOS软件时，偶然发现里面用到AES，但从伪代码看此AES算法和印象中不太一样，且对一些数据做了魔改。经过查找发现用到了TinyAES这个工程https://github.com/kokke/tiny-AES-c。但本人用Python顺手，就想办法用Python来测试之，于是有了本篇

开发

  参照https://www.0xaa55.com/thread-26563-1-1.html准备好pybind。观察TinyAes下的aes.h，编写代码

// aes.h中以下函数是我关心的，需要转成python使用。至于为什么不直接弄成一个aes_enc/aes_dec的函数而是要分成这三个函数，因为我还要测试魔改，便于接口扩展
void AES_init_ctx_iv(struct AES_ctx* ctx, const uint8_t* key, const uint8_t* iv);
void AES_CBC_encrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, size_t length);
void AES_CBC_decrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, size_t length);

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

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

PYBIND11_MODULE(pyutil, m) {
    m.def("tinyaes_aes_init_ctx_iv", [](const std::string& key, const std::string& iv) -> uint64_t {
        void* buf = malloc(sizeof(AES_ctx));
        memset(buf, 0, sizeof(AES_ctx));
        AES_ctx* ctx = (AES_ctx*)buf;
        AES_init_ctx_iv(ctx, (uint8_t*)key.data(), (uint8_t*)iv.data());
        return (uint64_t)(uintptr_t)buf;
    });
    m.def("tinyaes_aes_cbc_encrypt_buffer", [](uint64_t ibuf, const std::string& buf) -> py::bytes {
        std::string outbuf = buf;
        AES_ctx* ctx = (AES_ctx*)(uintptr_t)ibuf;
        AES_CBC_encrypt_buffer(ctx, (uint8_t*)outbuf.data(), outbuf.length());
        return py::bytes(outbuf.data(), outbuf.size());
    });
    m.def("tinyaes_aes_cbc_decrypt_buffer", [](uint64_t ibuf, const std::string& buf) -> py::bytes {
        std::string outbuf = buf;
        AES_ctx* ctx = (AES_ctx*)(uintptr_t)ibuf;
        AES_CBC_decrypt_buffer(ctx, (uint8_t*)outbuf.data(), outbuf.length());
        return py::bytes(outbuf.data(), outbuf.size());
    });
    m.def("tinyaes_aes_free", [](uint64_t ibuf) -> void {
        if (ibuf != 0) {
            free((void*)(uintptr_t)ibuf);
        }
    });
}

注意上述代码有以下技巧:

• aes.c是c代码，需要单独编译，且引入时需要extern "C"一下，且用gcc单独编译
• 因为Python无法处理c/c++指针，所以指针类型直接转换为uint64_t类型
• 如果返回类型为std::string则pybind返回给python的是str类型，而py::bytes返回的是bytes类型

编译模块:

#!/bin/bash
gcc -O3 -Wall -c -o aes.o aes.c
c++ -O3 -Wall -shared -std=c++11 -undefined dynamic_lookup $(python3 -m pybind11 --includes) \
    -DPYBIND11 pyutil.cpp aes.o -o pyutil$(python3-config --extension-suffix)

测试

依赖pycrypto

#! /usr/bin/env python3
# # -*- coding: utf-8 -*-

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"

import pyutil
ctx = pyutil.tinyaes_aes_init_ctx_iv(key, iv)
data1 = pyutil.tinyaes_aes_cbc_encrypt_buffer(ctx, data)
pyutil.tinyaes_aes_free(ctx)
print(data1.hex())

from Crypto.Cipher import AES
cipher = AES.new(key, AES.MODE_CBC, IV=iv)
data2 = cipher.encrypt(data)
print(data2.hex())

结果: 66e94bd4ef8a2c3b884cfa59ca342b2e

20230418改进

此次将aes处理为对象，可与上述方式做对比

#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

0xAA55

虽然但是 我比较倾向于使用 import ctypes 来调用 C/C++ 的 Dynamic Link Lib

lichao

0xAA55 发表于 2023-3-22 10:38
虽然但是 我比较倾向于使用 import ctypes 来调用 C/C++ 的 Dynamic Link Lib

此例确实可以用ctype。但ctype无法处理c++，也不适用于非导出函数，复杂的仍然要写wrapper
如果此例是个c++的源码，则可直接使用pybind，不用我这么转换了
各有各的优势，ctype/libffi这种方式如果有动态库直接就用了；而我这种方式是要一起编译成一个独立二进制，需要用到内部函数且可能需要对他源码进行魔改，一定要重编译