zoukankan      html  css  js  c++  java
  • 实战经验分享:使用 PyO3 来构建你的 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
    	
    rld";
            // 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 快速入门

  • 相关阅读:
    Mvc form提交
    FlexiGrid 使用 全选、自动绑定
    Mysql Insert Or Update语法例子
    orderby与groupby同时使用
    SQLSTATE[HY000] [2002] No such file or directory in
    swoole安装
    关于商城分类查询表结构与数据查询
    查询数据库每张表的信息
    php 正则验证
    PHP代码优化
  • 原文地址:https://www.cnblogs.com/upyun/p/15406555.html
Copyright © 2011-2022 走看看