ppaass_crypto/crypto/
rsa.rs

1use std::{fmt::Debug, path::Path, sync::Arc};
2use std::{fs, io::Read};
3
4use crate::error::CryptoError;
5use rand::rngs::OsRng;
6use rsa::{
7    pkcs8::{DecodePrivateKey, DecodePublicKey, EncodePrivateKey, EncodePublicKey, LineEnding},
8    Pkcs1v15Encrypt,
9};
10use rsa::{RsaPrivateKey, RsaPublicKey};
11
12const DEFAULT_AGENT_PRIVATE_KEY_PATH: &str = "AgentPrivateKey.pem";
13const DEFAULT_AGENT_PUBLIC_KEY_PATH: &str = "AgentPublicKey.pem";
14const DEFAULT_PROXY_PRIVATE_KEY_PATH: &str = "ProxyPrivateKey.pem";
15const DEFAULT_PROXY_PUBLIC_KEY_PATH: &str = "ProxyPublicKey.pem";
16
17/// The rsa crypto fetcher,
18/// each player have a rsa crypto
19/// which can be fund from the storage
20/// with user token
21pub trait RsaCryptoFetcher {
22    /// Fetch the rsa crypto by user token
23    fn fetch(&self, user_token: impl AsRef<str>) -> Result<Option<&RsaCrypto>, CryptoError>;
24}
25
26impl<T> RsaCryptoFetcher for Arc<T>
27where
28    T: RsaCryptoFetcher,
29{
30    fn fetch(&self, user_token: impl AsRef<str>) -> Result<Option<&RsaCrypto>, CryptoError> {
31        RsaCryptoFetcher::fetch(self.as_ref(), user_token)
32    }
33}
34
35impl<T> RsaCryptoFetcher for &T
36where
37    T: RsaCryptoFetcher,
38{
39    fn fetch(&self, user_token: impl AsRef<str>) -> Result<Option<&RsaCrypto>, CryptoError> {
40        RsaCryptoFetcher::fetch(*self, user_token)
41    }
42}
43
44impl<T> RsaCryptoFetcher for &mut T
45where
46    T: RsaCryptoFetcher,
47{
48    fn fetch(&self, user_token: impl AsRef<str>) -> Result<Option<&RsaCrypto>, CryptoError> {
49        RsaCryptoFetcher::fetch(*self, user_token)
50    }
51}
52
53/// The util to do RSA encryption and decryption.
54#[derive(Debug)]
55pub struct RsaCrypto {
56    /// The private used to do decryption
57    private_key: RsaPrivateKey,
58    /// The public used to do encryption
59    public_key: RsaPublicKey,
60}
61
62impl RsaCrypto {
63    pub fn new<A, B>(mut public_key_read: A, mut private_key_read: B) -> Result<Self, CryptoError>
64    where
65        A: Read + Debug,
66        B: Read + Debug,
67    {
68        let mut public_key_string = String::new();
69        public_key_read.read_to_string(&mut public_key_string)?;
70        let public_key = RsaPublicKey::from_public_key_pem(&public_key_string)
71            .map_err(|e| CryptoError::Other(format!("{e:?}")))?;
72        let mut private_key_string = String::new();
73        private_key_read.read_to_string(&mut private_key_string)?;
74        let private_key = RsaPrivateKey::from_pkcs8_pem(&private_key_string)
75            .map_err(|e| CryptoError::Other(format!("{e:?}")))?;
76        Ok(Self {
77            public_key,
78            private_key,
79        })
80    }
81
82    pub fn encrypt(&self, target: &[u8]) -> Result<Vec<u8>, CryptoError> {
83        let result = self
84            .public_key
85            .encrypt(&mut OsRng, Pkcs1v15Encrypt, target.as_ref())
86            .map_err(|e| CryptoError::Other(format!("{e:?}")))?;
87        Ok(result)
88    }
89
90    pub fn decrypt(&self, target: &[u8]) -> Result<Vec<u8>, CryptoError> {
91        let result = self
92            .private_key
93            .decrypt(Pkcs1v15Encrypt, target.as_ref())
94            .map_err(|e| CryptoError::Other(format!("{e:?}")))?;
95        Ok(result)
96    }
97}
98
99pub fn generate_agent_key_pairs(base_dir: &str, user_token: &str) -> Result<(), CryptoError> {
100    let private_key_path = format!("{base_dir}/{user_token}/{DEFAULT_AGENT_PRIVATE_KEY_PATH}");
101    let private_key_path = Path::new(private_key_path.as_str());
102    let public_key_path = format!("{base_dir}/{user_token}/{DEFAULT_AGENT_PUBLIC_KEY_PATH}");
103    let public_key_path = Path::new(public_key_path.as_str());
104    generate_rsa_key_pairs(private_key_path, public_key_path)
105}
106
107pub fn generate_proxy_key_pairs(base_dir: &str, user_token: &str) -> Result<(), CryptoError> {
108    let private_key_path = format!("{base_dir}/{user_token}/{DEFAULT_PROXY_PRIVATE_KEY_PATH}");
109    let private_key_path = Path::new(private_key_path.as_str());
110    let public_key_path = format!("{base_dir}/{user_token}/{DEFAULT_PROXY_PUBLIC_KEY_PATH}");
111    let public_key_path = Path::new(public_key_path.as_str());
112    generate_rsa_key_pairs(private_key_path, public_key_path)
113}
114
115fn generate_rsa_key_pairs(
116    private_key_path: &Path,
117    public_key_path: &Path,
118) -> Result<(), CryptoError> {
119    let private_key = RsaPrivateKey::new(&mut OsRng, 2048).expect("Fail to generate private key");
120    let public_key = RsaPublicKey::from(&private_key);
121    let private_key_pem = private_key
122        .to_pkcs8_pem(LineEnding::CRLF)
123        .expect("Fail to generate pem for private key.");
124    let public_key_pem = public_key
125        .to_public_key_pem(LineEnding::CRLF)
126        .expect("Fail to generate pem for public key.");
127    match private_key_path.parent() {
128        None => {
129            println!("Write private key: {:?}", private_key_path.to_str());
130            fs::write(private_key_path, private_key_pem.as_bytes())?;
131        }
132        Some(parent) => {
133            if !parent.exists() {
134                println!("Create parent directory :{:?}", parent.to_str());
135                fs::create_dir_all(parent)?;
136            }
137            println!("Write private key: {:?}", private_key_path.to_str());
138            fs::write(private_key_path, private_key_pem.as_bytes())?;
139        }
140    };
141    match public_key_path.parent() {
142        None => {
143            println!("Write public key: {:?}", public_key_path.to_str());
144            fs::write(public_key_path, public_key_pem.as_bytes())?;
145        }
146        Some(parent) => {
147            if !parent.exists() {
148                println!("Create parent directory :{:?}", parent.to_str());
149                fs::create_dir_all(parent)?;
150            }
151            println!("Write public key: {:?}", public_key_path.to_str());
152            fs::write(public_key_path, public_key_pem.as_bytes())?;
153        }
154    };
155    Ok(())
156}
ppaass_crypto/crypto/rsa.rs

ppaass_crypto/crypto/
rsa.rs
