9

实战经验分享:使用 PyO3 来构建你的 Python 模块

 3 years ago
source link: https://segmentfault.com/a/1190000040811460
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

实战经验分享:使用 PyO3 来构建你的 Python 模块

发布于 10 月 14 日

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/pyo...

后续又扩展了一些内容,比如 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 快速入门


Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK