shell_gpt/
config.rs

1use crate::errors::ShellGptError;
2use crate::{encryption, openai};
3use anyhow::{anyhow, Context};
4use directories::ProjectDirs;
5use std::fs::File;
6use std::io::prelude::*;
7use std::path::PathBuf;
8use std::{env, fs, io};
9use users::{get_current_uid, get_user_by_uid};
10
11#[derive(Debug)]
12pub struct CliArgs {
13    pub input: String,
14    pub show_help: bool,
15    pub clear_saved_config: bool,
16    pub pre_prompt: openai::PrePrompt,
17    /// Just print the answer - No spinner, pretty-print, or interactive action
18    pub raw: bool,
19}
20
21#[derive(Debug)]
22pub struct Config {
23    pub openai_api_key: String,
24}
25
26impl Config {
27    pub fn load_config(_cli_args: &CliArgs) -> Config {
28        let openai_api_key = load_api_key()
29            .unwrap_or_else(|_| register_api_key().unwrap_or_else(|e| panic!("{:#?}", e)));
30        Config { openai_api_key }
31    }
32
33    pub fn clear_saved_config() {
34        let path = get_config_dir_path();
35        fs::remove_dir_all(path).unwrap_or_else(|e| panic!("{:#?}", e));
36        println!("Local configuration was cleared!");
37    }
38
39    pub fn parse_cli_args(args: Vec<String>) -> CliArgs {
40        let args: Vec<String> = args.iter().skip(1).cloned().collect();
41
42        let input = args.join(" ");
43        let mut show_help = false;
44        let mut clear_saved_config = false;
45        let mut pre_prompt = openai::PrePrompt::NoPrePrompt;
46        let mut raw = false;
47
48        args.iter().for_each(|x| match x.trim() {
49            "--help" | "-h" => show_help = true,
50            "--shell" | "--bash" | "--script" | "-s" => pre_prompt = openai::PrePrompt::ShellScript,
51            "--remove-config" | "--delete-config" | "--clear-config" => clear_saved_config = true,
52            "--raw" => raw = true,
53            _ => {}
54        });
55
56        CliArgs {
57            input,
58            show_help,
59            clear_saved_config,
60            pre_prompt,
61            raw,
62        }
63    }
64}
65
66fn load_api_key() -> anyhow::Result<String> {
67    if let Ok(api_key) = read_api_key_from_env_var() {
68        return Ok(api_key);
69    }
70
71    let config_path = get_config_openai_api_key_path();
72    let mut file =
73        File::open(&config_path).with_context(|| format!("Failed to open file {config_path:?}"))?;
74    let mut buffer = Vec::new();
75    file.read_to_end(&mut buffer)
76        .with_context(|| format!("Failed to read file {config_path:?}"))?;
77
78    let api_key_encrypted = encryption::decrypt(buffer, get_config_encryption_password())?;
79
80    let api_key = String::from_utf8(api_key_encrypted)
81        .with_context(|| "Failed to convert decrypted API key to UTF-8")?;
82    key_not_empty(api_key)
83}
84
85fn save_api_key(api_key: &str) -> anyhow::Result<()> {
86    let config_path = get_config_openai_api_key_path();
87    println!("config_path = {config_path:?}");
88
89    // Basically `mkdir -p /home/<user>/.config/shell-gpt-rs`
90    let prefix = config_path.parent().unwrap();
91    fs::create_dir_all(prefix)?;
92
93    // Write encrypted API key
94    let mut file = File::create(&config_path)?;
95    let api_key_encrypted =
96        encryption::encrypt(api_key.as_bytes(), get_config_encryption_password())?;
97
98    file.write_all(&*api_key_encrypted)
99        .with_context(|| format!("Could not save the API key to {config_path:?}"))
100}
101
102pub fn register_api_key() -> anyhow::Result<String> {
103    println!("{}", get_config_dir_path().to_str().unwrap().to_string());
104    println!("You need to enter an OpenAI API key");
105    println!();
106    println!("You can create a new API key at https://platform.openai.com/account/api-keys");
107    println!("Your key will be encrypted and saved locally for future use");
108    println!();
109    println!("If you don't want your key to be saved, you can pass it using the `OPENAI_KEY` environment variable");
110    println!();
111    println!();
112    println!("Enter your OpenAI API key: ");
113
114    let mut api_key = String::new();
115
116    io::stdin()
117        .read_line(&mut api_key)
118        .context("Failed to read your input")?;
119
120    match api_key.trim() {
121        api_key if !api_key.is_empty() => {
122            save_api_key(&api_key).context("Failed saving the API key")?;
123            // save_api_key(&api_key).context("Failed saving the API key")?;
124            Ok(api_key.to_string())
125        }
126        e => Err(anyhow!("Received an empty API key - {}", e)),
127    }
128}
129
130fn read_api_key_from_env_var() -> anyhow::Result<String> {
131    match env::var("OPENAI_KEY") {
132        Ok(key) => key_not_empty(key),
133        Err(_) => Err(ShellGptError::ApiKeyMissing.into()),
134    }
135}
136
137fn key_not_empty(key: String) -> anyhow::Result<String> {
138    if !key.trim().is_empty() {
139        Ok(key.trim().to_string())
140    } else {
141        Err(ShellGptError::ApiKeyEmpty.into())
142    }
143}
144
145fn get_config_dir_path<'a>() -> PathBuf {
146    ProjectDirs::from("com", "shell-gpt-rs", "shell-gpt-rs")
147        .unwrap()
148        .config_dir()
149        .to_path_buf()
150}
151
152fn get_config_openai_api_key_path() -> PathBuf {
153    get_config_dir_path().join("openai_api_key.encrypted.txt")
154}
155
156/// Generate the password used to encrypt the configuration file with.
157///
158/// Password is `username_ENCRYPTION_PASSWORD_SUFFIX`
159fn get_config_encryption_password() -> String {
160    const ENCRYPTION_PASSWORD_SUFFIX: &'static str =
161        "shell-gpt-rs-EL9Kaesj7Q6pc9BzsfxVpjPbNnuj8bGJ";
162    let user = get_user_by_uid(get_current_uid()).unwrap();
163    format!(
164        "{}_{}",
165        user.name().to_str().unwrap(),
166        ENCRYPTION_PASSWORD_SUFFIX
167    )
168}
shell_gpt/config.rs

shell_gpt/
config.rs
