實戰經驗分享：使用 PyO3 來構建你的 Python 模組

雲叔_又拍雲發表於2021-10-14
Python

PyO3 主要用於建立原生 Python 的擴充套件模組。PyO3 還支援從 Rust 二進位制檔案執行 Python 程式碼並與之互動，可以實現 rust 與 Python 程式碼共存。在一些對效能要求較高的模組上，可以考慮使用 PyO3 構建對應的功能模組。PyO3 的功能分離，不用過多擔心模組之間的耦合性，並且在速度上能有一定的提升。

github地址: https://github.com/PyO3/pyo3

版本規定如下：

  • Python 3.6+

  • Rust 1.41+

接下來我們通過一個小的 demo 瞭解一下從 PyO3 編譯模組到 Python 中正常使用的整個流程。

cargo new --lib string-sum

建立專案

# lib.rs
[package]
name = "string-sum"
version = "0.1.0"
edition = "2018"

[lib]
name = "string_sum"
# "cdylib" is necessary to produce a shared library for Python to import from.
#
# Downstream Rust code (including code in `bin/`, `examples/`, and `tests/`) will not be able
# to `use string_sum;` unless the "rlib" or "lib" crate type is also included, e.g.:
# crate-type = ["cdylib", "rlib"]
crate-type = ["cdylib"]

[dependencies.pyo3]
version = "0.14.1"
features = ["extension-module"] // 擴充套件模組，像其他的還有auto-initialize

// src/lib.rs
use std::usize;

use  pyo3::prelude::*;

// like this
// def sum_as_string(a:str, b:str) -> str:
//      return a+b
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String>{
    Ok((a+b).to_string())
}

// Mount method to module 
#[pymodule]
fn string_sum(py: Python, m: &PyModule) -> PyResult<()>{
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

編譯與使用

編譯完成之後，我們會在 target 資料夾下面發現一個 wheel 檔案。檔名組合為 “模組名 + 當前 Python 版本+當前系統型號”，比如：string_sum-0.1.0-cp39-cp39-macosx_10_7_x86_64.whl


pip3 install ./target/wheel/string_sum-0.1.0-cp39-cp39-macosx_10_7_x86_64.whl

建立 python 檔案：

# example.py
from string_sum import sum_as_string
print(sum_as_string(1,2))
# echo 3

編譯工具的選擇和使用

官方提供了兩種編譯工具的選擇：

  • rust 寫的 maturin

  • 傳統的setup.py的方式

使用 maturin 編譯

# 安裝 
pip3 install maturin
# 編譯
maturin build
# maturin publish 釋出
# 虛擬環境中使用 會自動去尋找/target/wheel/ 下的 *.wheel檔案然後安裝
virtualenv venv
source ./venv/bin/activate
maturin develop

使用 setup.py 編譯

# 安裝
pip3 install setuptools-rust

編寫 setup.py 檔案：

# setup.py


from setuptools import setup
from setuptools_rust import Binding, RustExtension

setup(
    # 包名稱
    name="string_sum", 
    # 包版本 
    version="0.1",
    # rust擴充套件 其中"string_sum.string_sum"中
    # 第一個string_sum 指的是當前的包
    # 第二個指的是
    # #[pymodule]
    # fn string_sum(py: Python, m: &PyModule) -> PyResult<()>{
    #     m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    #     Ok(())
    # }
    # 中的string_sum
    rust_extensions=[
        RustExtension(
            "string_sum.string_sum", 
            binding=Binding.PyO3,
            debug=False
            )
    ],
    # 需要建立一個資料夾 string_sum
    packages=["string_sum"],
    # rust extensions are not zip safe, just like C-extensions.
    zip_safe=False,
    # 標註
    classifiers=[
        "License :: OSI Approved :: MIT License",
        "Development Status :: 3 - Alpha",
        "Intended Audience :: Developers",
        "Programming Language :: Python",
        "Programming Language :: Rust",
        "Operating System :: POSIX",
        "Operating System :: MacOS :: MacOS X",
    ],
    include_package_data=True
)

# 打包
mkdir string_sum
touch string_sum/__init__.py
virtualenv venv && source venv/bin/activate
pip setup.py build && pip setup.py install && pip setup.py develop

會引用本地的檔案：

docker 中的應用

同樣的，如果建立的 App 本身是在 docker 內部執行的。那麼第一步我們需要安裝 rust 的環境 dockerfile。具體如下：

#!/bin/bash
curl https://sh.rustup.rs -sSf | bash -s -- -y
source $HOME/.cargo/env
rustc --version
python setup.py install

# ddockerfile 
FROM python:3.7
WORKDIR /app
ADD . /app
RUN pip install --upgrade pip \
    && pip install -r requirements.txt
RUN ./init.sh
CMD [python, xx.py]

# requirements.txt
semantic-version==2.8.5
setuptools-rust==0.12.1
toml==0.10.2

# rust國內映象源 config
# /root/.cargo/config
[source.crates-io]
registry = "https://github.com/rust-lang/crates.io-index"
replace-with = 'ustc'
[source.ustc]
registry = "git://mirrors.ustc.edu.cn/crates.io-index"
[term]
verbose = true
color = 'auto'

具體目錄如下：

-rw-r--r-- Cargo.lock
-rw-r--r-- Cargo.toml
-rw-r--r-- config           # 配置檔案
-rw-r--r-- Dockerfile
-rwxrwxrwx init.sh          # 初始化rust環境指令碼
-rw-r--r-- requirements.txt
-rw-r--r-- setup.py         # 打包指令碼
drwxr-xr-x src              # rust專案
drwxr-xr-x string_sum 
-rw-r--r-- xx.py            # 可行性測試檔案

用 PyO3 寫一個 Python 的rsa加解密包

看過之前的文章的小夥伴《靈魂畫手：漫畫圖解 SSH》 ，應該對 rsa 的整個加解密流程有所瞭解啦。那我們不妨用 PyO3 來構建一個 Python 的 rsa 加解密包。使用場景如下：

客戶端本地生成公私鑰，通過前期認證過程，將公鑰傳送給服務端儲存，後期通訊過程中,客戶端主動傳送訊息給服務端，客戶端通過私鑰對資訊加密，服務端通過對應的公鑰進行解密。

github 地址: https://github.com/hzjsea/pyo3-crypto

後續又擴充套件了一些內容，比如 MD5 加密，簽名等等。

# 自動化指令碼
#!/bin/bash

echo "init......"

# set python version 
# INSTALL_PYTHON_VERSION=python3.6

find_python() {
        set +e
        unset BEST_VERSION
        for V in 37 3.7 38 3.8 39 3.9 3; do
                if which python$V >/dev/null; then
                        if [ "$BEST_VERSION" = "" ]; then
                                BEST_VERSION=$V
                        fi
                fi
        done
        echo $BEST_VERSION
        set -e
}

if [ "$INSTALL_PYTHON_VERSION" = "" ]; then
        INSTALL_PYTHON_VERSION=$(find_python)
fi

# This fancy syntax sets INSTALL_PYTHON_PATH to "python3.7", unless
# INSTALL_PYTHON_VERSION is defined.
# If INSTALL_PYTHON_VERSION equals 3.8, then INSTALL_PYTHON_PATH becomes python3.8
# 找不到就python3.7
INSTALL_PYTHON_PATH=python${INSTALL_PYTHON_VERSION:-3.7}
echo $INSTALL_PYTHON_PATH

echo "Python version is $INSTALL_PYTHON_VERSION"
$INSTALL_PYTHON_PATH -m venv venv
if [ ! -f "activate" ]; then
        ln -s venv/bin/activate .
fi

. ./activate

python -m pip install --upgrade pip
python -m pip install wheel
python -m pip install -r ./requirements.txt
maturin build
maturin develop

current_shell=$(echo $SHELL)
if current_shell=/bin/bash; then
    echo  "PASS: source /venv/bin/activate >> ~/.bashrc"
elif current_shell=/bin/zsh;then
    echo "PASS: source /venv/bin/activate >> ~/.zshrc"
fi

//  src/lib.rs 檔案
use std::u32;
use pyo3::prelude::*;
use openssl::rsa::{Padding,Rsa};

const SECRET: &'static str = "CHFfxQA3tqEZgKusgwZjmI5lFsoZxXGXnQLA97oYga2M33sLwREZyy1mWCM8GIIA";

mod crypto_utils {
    use hmac::{Hmac, Mac, NewMac};
    use sha2::Sha256;
    use std::fmt::Write;

    type Hmacsha256 = Hmac<Sha256>;
    fn encode_hex(bytes: &[u8]) -> String {
        let mut s = String::with_capacity(bytes.len() * 2);
        for &b in bytes {
            match write!(&mut s, "{:02x}", b) {
                Ok(_) => {},
                Err(_) => {}
            };
        }
        s
    }

    pub fn hash_hmac(secret: &str, msg: &str) -> String {
        let mut mac = Hmacsha256::new_from_slice(secret.as_bytes()).expect("HMAC can take key of any size");
        mac.update(msg.as_bytes());
        let result = mac.finalize();
        let code_bytes = result.into_bytes();
        encode_hex(&code_bytes)
    }

}

// create public/private key  create_key(1024)
fn create_key(len:u32) -> (String,String){
    let rsa = openssl::rsa::Rsa::generate(len).unwrap();
    let pubkey = String::from_utf8(rsa.public_key_to_pem().unwrap()).unwrap();
    let prikey  = String::from_utf8(rsa.private_key_to_pem().unwrap()).unwrap();

    (pubkey, prikey)
}



#[pyclass]
struct Crypto {
    // #[pyo3(get, set)]
    // pubkey: String,
    // #[pyo3(get,set)]
    // prikey: String,
    pub_key: Rsa<openssl::pkey::Public>,
    pri_key: Rsa<openssl::pkey::Private>
}

#[pyfunction]
fn generate_key(len:u32) -> (String, String){
    create_key(len)
}

#[pymethods]
impl Crypto {
    #[new]
    pub fn __new__(pubkey: &str,prikey: &str) -> Self {
        Crypto {
            // pubkey: pubkey.to_owned(),
            // prikey: prikey.to_owned(),
            pub_key: Rsa::public_key_from_pem(pubkey.as_bytes()).unwrap(),
            pri_key: Rsa::private_key_from_pem(prikey.as_bytes()).unwrap(),
        }
    }

    // public decrypt 
    pub fn public_decrypt(&self, msg:&str) -> String {
        let mut out: [u8; 4096] = [0;4096];
        let decoded = openssl::base64::decode_block(msg).unwrap();
        if let Ok(size) = self.pub_key.public_decrypt(&decoded, &mut out, Padding::PKCS1) {
            let real_size = if size > 4096 {4096} else {size};
            // openssl::base64::encode_block(&out[..real_size])
            String::from_utf8(out[..real_size].to_vec()).unwrap()
        } else {
            String::default()
        }
    }

    // public encrypt 
    pub fn public_encrypt(&self, msg:&str) -> String {
        let mut out: [u8; 4096] = [0;4096];
        if let Ok(size) = self.pub_key.public_encrypt(msg.as_bytes(), &mut out, Padding::PKCS1) {
            let real_size = if size > 4096 {4096}else{size};
            openssl::base64::encode_block(&out[..real_size])
        } else {
            String::default()
        }
    }

    // private encrypt
    pub fn private_encrypt(&self, msg:&str) -> String{
        let mut out: [u8; 4096] = [0;4096];
        if let Ok(size) = self.pri_key.private_encrypt(msg.as_bytes(), &mut out, Padding::PKCS1) {
            let real_size = if size > 4096 {4096}else{size};
            openssl::base64::encode_block(&out[..real_size])
        } else {
            String::default()
        }
    }

    // private decrypt
    pub fn private_decrypt(&self, msg: &str) -> String{
        let mut out: [u8; 4096] = [0;4096];
        let decoded = openssl::base64::decode_block(msg).unwrap();
        if let Ok(size) = self.pri_key.private_decrypt(&decoded, &mut out, Padding::PKCS1) {
            let real_size = if size > 4096 {4096} else {size};
            // openssl::base64::encode_block(&out[..real_size])
            String::from_utf8(out[..real_size].to_vec()).unwrap()
        } else {
            String::default()
        } 
    }

    // sign
    pub fn sign(&self, msg: &str) ->String {
        crypto_utils::hash_hmac(SECRET, msg)
    }
}

#[pymodule]
fn yacrypto(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Crypto>()?;
    m.add_function(wrap_pyfunction!(generate_key, m)?).unwrap();
    Ok(())
}


#[cfg(test)]
mod tests {
    use base64;
    #[test]
    fn works(){

        // create rsa
        let rsa = openssl::rsa::Rsa::generate(1024).unwrap();
        // create public key 
        let public_key = rsa.public_key_to_pem().unwrap();
        println!("{:?}", String::from_utf8(public_key.clone()));
        let private_key = rsa.private_key_to_pem().unwrap();


        let data = "hellowo\n\t\rrld";
        // public encrypt 
        let mut buf:Vec<u8> = vec![0;rsa.size() as usize];
        let rsa_pub = openssl::rsa::Rsa::public_key_from_pem(&public_key).unwrap();
        let _ = rsa_pub.public_encrypt(data.as_bytes(), &mut buf , openssl::rsa::Padding::PKCS1);

        // private decrypt => 
        let data = buf;
        let mut buf:Vec<u8> = vec![0;rsa.size() as usize];
        let rsa_pri  = openssl::rsa::Rsa::private_key_from_pem(&private_key).unwrap();
        if let Ok(size) = rsa_pri.private_decrypt(&data, &mut buf, openssl::rsa::Padding::PKCS1){
            let real_size = if size > 1024 {1024} else {size};
            let buf = &buf[..real_size];
            let base64_string = openssl::base64::encode_block(&buf);
            let result = base64::decode(base64_string);
            println!("{:?}",result);
            // println!("buf => {:?}",openssl::base64::encode_block(&buf))

            let echo_str = String::from_utf8((&buf).to_vec()).unwrap();
            println!("{:?}",echo_str);

        }
    }
}

推薦閱讀

webpack 從 0 到 1 構建 vue

Ansible 快速入門

相關文章