ppaass_v3_proxy_tool/handler/
generate_user.rs

1use crate::config::ProxyToolConfig;
2use crate::crypto::{generate_agent_key_pairs, generate_proxy_key_pairs};
3use anyhow::{Result, anyhow};
4use chrono::{TimeDelta, Utc};
5use ppaass_common::crypto::{
6    DEFAULT_AGENT_PRIVATE_KEY_PATH, DEFAULT_AGENT_PUBLIC_KEY_PATH, DEFAULT_PROXY_PRIVATE_KEY_PATH,
7    DEFAULT_PROXY_PUBLIC_KEY_PATH,
8};
9use ppaass_common::generate_uuid;
10use ppaass_common::user::repo::fs::{
11    FS_USER_INFO_CONFIG_FILE_NAME, FsAgentUserInfoContent, FsProxyUserInfoContent,
12};
13use std::io::Write;
14use std::net::SocketAddr;
15use std::ops::Add;
16use std::path::{Path, PathBuf};
17use std::str::FromStr;
18use zip::write::SimpleFileOptions;
19const DEFAULT_SEND_TO_AGENT_DIR: &str = "send_to_agent";
20const DEFAULT_TEMP_DIR: &str = "temp";
21
22pub struct GenerateUserHandlerArgument {
23    pub username: String,
24    pub temp_dir: Option<PathBuf>,
25    pub agent_rsa_dir: Option<PathBuf>,
26    pub expire_after_days: Option<i64>,
27    pub proxy_servers: Vec<String>,
28}
29
30fn generate_agent(temp_user_dir: &Path, arg: &GenerateUserHandlerArgument) -> Result<()> {
31    println!(
32        "Begin to generate agent user info configuration file for: {}",
33        &arg.username
34    );
35    let default_send_to_agent_dir = PathBuf::from(DEFAULT_SEND_TO_AGENT_DIR);
36    if !default_send_to_agent_dir.exists() {
37        std::fs::create_dir_all(&default_send_to_agent_dir)?;
38    }
39    let proxy_public_key_file_content = std::fs::read_to_string(
40        temp_user_dir
41            .join(&arg.username)
42            .join(DEFAULT_PROXY_PUBLIC_KEY_PATH),
43    )?;
44    let agent_private_key_file_content = std::fs::read_to_string(
45        temp_user_dir
46            .join(&arg.username)
47            .join(DEFAULT_AGENT_PRIVATE_KEY_PATH),
48    )?;
49
50    if arg.proxy_servers.is_empty() {
51        eprintln!("Proxy servers should not be empty.");
52        return Err(anyhow!("Proxy servers should not be empty."));
53    }
54    for proxy_server in arg.proxy_servers.iter() {
55        if let Err(e) = SocketAddr::from_str(&proxy_server) {
56            eprintln!("Failed to parse proxy server: {proxy_server}");
57            return Err(anyhow!("Fail to parse proxy server address: {e:?}"));
58        }
59    }
60
61    let agent_user_info = FsAgentUserInfoContent::new(
62        arg.proxy_servers.clone(),
63        DEFAULT_PROXY_PUBLIC_KEY_PATH.to_string(),
64        DEFAULT_AGENT_PRIVATE_KEY_PATH.to_string(),
65    );
66    let agent_user_info_config_file_content = toml::to_string(&agent_user_info)?;
67    let zip_file_path = default_send_to_agent_dir.join(format!("{}.zip", &arg.username));
68    let zip_file_path = Path::new(&zip_file_path);
69    let zip_file = std::fs::File::create(zip_file_path)?;
70    let mut zip_file_writer = zip::ZipWriter::new(zip_file);
71    let zi_file_options =
72        SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
73    zip_file_writer.start_file(FS_USER_INFO_CONFIG_FILE_NAME, zi_file_options)?;
74    zip_file_writer.write_all(agent_user_info_config_file_content.as_bytes())?;
75    zip_file_writer.start_file(DEFAULT_PROXY_PUBLIC_KEY_PATH, zi_file_options)?;
76    zip_file_writer.write_all(proxy_public_key_file_content.as_bytes())?;
77    zip_file_writer.start_file(DEFAULT_AGENT_PRIVATE_KEY_PATH, zi_file_options)?;
78    zip_file_writer.write_all(agent_private_key_file_content.as_bytes())?;
79    println!("Success write agent user configuration to: {zip_file_path:?}",);
80    Ok(())
81}
82
83pub fn generate_proxy(
84    config: &ProxyToolConfig,
85    temp_user_dir: &Path,
86    arg: &GenerateUserHandlerArgument,
87) -> Result<()> {
88    println!(
89        "Begin to generate proxy user info configuration file for: {}",
90        &arg.username
91    );
92    let proxy_user_dir = config.user_dir();
93    if !proxy_user_dir.exists() {
94        std::fs::create_dir_all(&proxy_user_dir)?;
95    }
96    let proxy_private_key_file_content = std::fs::read_to_string(
97        temp_user_dir
98            .join(&arg.username)
99            .join(DEFAULT_PROXY_PRIVATE_KEY_PATH),
100    )?;
101
102    let agent_public_key_file_content = std::fs::read_to_string(
103        temp_user_dir
104            .join(&arg.username)
105            .join(DEFAULT_AGENT_PUBLIC_KEY_PATH),
106    )?;
107
108    let expired_date_time = match arg.expire_after_days {
109        None => None,
110        Some(days) => Some(Utc::now().add(TimeDelta::days(days))),
111    };
112
113    let proxy_user_info = FsProxyUserInfoContent::new(
114        expired_date_time,
115        DEFAULT_AGENT_PUBLIC_KEY_PATH.to_string(),
116        DEFAULT_PROXY_PRIVATE_KEY_PATH.to_string(),
117    );
118    let proxy_user_info_config_file_content = toml::to_string(&proxy_user_info)?;
119
120    let zip_file_path = proxy_user_dir.join(format!("{}.zip", &arg.username));
121    let zip_file_path = Path::new(&zip_file_path);
122    let zip_file = std::fs::File::create(zip_file_path)?;
123    let mut zip_file_writer = zip::ZipWriter::new(zip_file);
124    let zi_file_options =
125        SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
126    zip_file_writer.start_file(FS_USER_INFO_CONFIG_FILE_NAME, zi_file_options)?;
127    zip_file_writer.write_all(proxy_user_info_config_file_content.as_bytes())?;
128    zip_file_writer.start_file(DEFAULT_PROXY_PRIVATE_KEY_PATH, zi_file_options)?;
129    zip_file_writer.write_all(proxy_private_key_file_content.as_bytes())?;
130    zip_file_writer.start_file(DEFAULT_AGENT_PUBLIC_KEY_PATH, zi_file_options)?;
131    zip_file_writer.write_all(agent_public_key_file_content.as_bytes())?;
132    println!("Success write proxy user configuration to: {zip_file_path:?}",);
133    Ok(())
134}
135
136pub fn generate_user(config: &ProxyToolConfig, arg: GenerateUserHandlerArgument) -> Result<()> {
137    let default_tmp_dir = PathBuf::from(DEFAULT_TEMP_DIR);
138    let temp_dir = arg.temp_dir.as_ref().unwrap_or_else(|| &default_tmp_dir);
139    let temp_user_dir = temp_dir.join(generate_uuid());
140    println!(
141        "Begin to generate RSA key for [{}] in [{:?}]",
142        arg.username, temp_user_dir
143    );
144    generate_proxy_key_pairs(&temp_user_dir, &arg.username)?;
145    generate_agent_key_pairs(&temp_user_dir, &arg.username)?;
146    generate_proxy(config, &temp_user_dir, &arg)?;
147    generate_agent(&temp_user_dir, &arg)?;
148    Ok(())
149}