了解PyO3

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/wheelring_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 .env/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

                    1.png

                    会引用本地的文件:

                    2.png

                    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 env/bin/activate >> ~/.bashrc"
                                elif current_shell=/bin/zsh;then
                                    echo "PASS: source env/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);
                                  
                                          }
                                      }
                                  }