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

lichao 发表于 2023-3-19 17:44:54

本帖最后由 lichao 于 2023-4-18 13:23 编辑

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

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

```c
// 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);
```

```cpp
#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类型

编译模块:

```bash
#!/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

```Python
#! /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处理为对象，可与上述方式做对比
```cpp
#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())
```

得到:

```txt
origin: 00000000000000000000000000000000
encrypt: 66e94bd4ef8a2c3b884cfa59ca342b2e
decrypt: 00000000000000000000000000000000
```

0xAA55 发表于 2023-3-22 10:38:54

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

lichao 发表于 2023-3-22 11:05:08

本帖最后由 lichao 于 2023-3-22 13:42 编辑

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

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

页: [1]

技术宅的结界's Archiver

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