sudo_gcp/
lib.rs

1use std::{
2    collections::HashSet,
3    fmt::Display,
4    process::{Command, Stdio},
5    str::FromStr,
6    time::Duration,
7};
8
9use chrono::{DateTime, Utc};
10use keyring::Entry;
11use reqwest::{blocking::Client, header::HeaderMap};
12use serde::{Deserialize, Serialize};
13
14const DEFAULT_OAUTH_SCOPES: &[&str] = &["https://www.googleapis.com/auth/cloud-platform"];
15
16const DEFAULT_LIFETIME_SECONDS: u64 = 3600;
17const IAM_API: &str = "https://iamcredentials.googleapis.com/v1";
18static USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
19
20#[derive(Debug, Serialize, Deserialize)]
21#[serde(rename_all = "camelCase")]
22pub struct AccessToken(String);
23
24impl FromStr for AccessToken {
25    type Err = String;
26
27    fn from_str(s: &str) -> Result<Self, Self::Err> {
28        Ok(Self(s.to_string()))
29    }
30}
31
32impl AsRef<str> for AccessToken {
33    fn as_ref(&self) -> &str {
34        self.0.as_ref()
35    }
36}
37
38impl From<String> for AccessToken {
39    fn from(value: String) -> Self {
40        Self(value)
41    }
42}
43
44#[derive(Debug)]
45pub struct GcloudConfig {
46    _account: String,
47    access_token: AccessToken,
48}
49
50impl FromStr for GcloudConfig {
51    type Err = String;
52
53    fn from_str(s: &str) -> Result<Self, Self::Err> {
54        let (account, access_token) = s.trim().split_once(',').expect("config-helper call failed");
55        Ok(Self {
56            _account: account.to_string(),
57            access_token: AccessToken::from_str(access_token)
58                .expect("failed to parse access token"),
59        })
60    }
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct Email(String);
65
66impl FromStr for Email {
67    type Err = String;
68
69    fn from_str(s: &str) -> Result<Self, Self::Err> {
70        Ok(Self(s.to_string()))
71    }
72}
73
74impl AsRef<str> for Email {
75    fn as_ref(&self) -> &str {
76        self.0.as_ref()
77    }
78}
79
80impl Display for Email {
81    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82        write!(f, "{}", self.0)
83    }
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
87pub struct Scopes(HashSet<String>);
88
89impl FromStr for Scopes {
90    type Err = String;
91
92    fn from_str(s: &str) -> Result<Self, Self::Err> {
93        let scopes = s.split(',').map(|s| s.to_string()).collect();
94        Ok(Self(scopes))
95    }
96}
97
98impl Display for Scopes {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        let sorted_scopes: Vec<String> = self.0.iter().map(|s| s.to_string()).collect();
101        let scopes: String = sorted_scopes.join(",");
102        write!(f, "{}", scopes)
103    }
104}
105impl Default for Scopes {
106    fn default() -> Self {
107        let owned_scopes: HashSet<String> = DEFAULT_OAUTH_SCOPES
108            .iter()
109            .map(|scope| scope.to_string())
110            .collect();
111        Self(owned_scopes)
112    }
113}
114
115impl Scopes {
116    pub fn append_scopes(&self, additional_scopes: Scopes) -> Self {
117        let mut scopes = Scopes::default();
118        scopes.0.extend(additional_scopes.0);
119        scopes
120    }
121}
122
123#[derive(Debug, Clone, Serialize, Deserialize)]
124pub struct Lifetime(u64);
125
126// impl Serialize for Lifetime {
127//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
128//     where
129//         S: serde::Serializer,
130//     {
131//         serializer.serialize_u64(self.0.as_secs())
132//     }
133// }
134
135// impl FromStr for Lifetime {
136//     type Err = String;
137
138//     fn from_str(s: &str) -> Result<Self, Self::Err> {
139//         let trimmed_s = s.trim_end_matches('s');
140//         let seconds: u64 = trimmed_s.parse::<u64>().expect("failed to convert number");
141//         Ok(Self(Duration::from_secs(seconds)))
142//     }
143// }
144
145// impl From<u64> for Lifetime {
146//     fn from(value: u64) -> Self {
147//         Self(Duration::from_secs(value))
148//     }
149// }
150
151impl Display for Lifetime {
152    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153        write!(f, "{}s", self.0)
154    }
155}
156
157impl Default for Lifetime {
158    fn default() -> Self {
159        Self(DEFAULT_LIFETIME_SECONDS)
160    }
161}
162
163pub fn get_gcloud_config() -> GcloudConfig {
164    let config = Command::new("gcloud")
165        .args([
166            "config",
167            "config-helper",
168            "--format",
169            "csv[no-heading](configuration.properties.core.account,credential.access_token)",
170        ])
171        .stderr(Stdio::inherit())
172        .output()
173        .expect("gcloud call failed");
174    GcloudConfig::from_str(std::str::from_utf8(&config.stdout).unwrap()).unwrap()
175}
176
177#[derive(Debug, Serialize, Deserialize, Default)]
178struct TokenRequest {
179    lifetime: String,
180    scope: Scopes,
181}
182
183#[derive(Debug, Deserialize, Serialize)]
184#[serde(rename_all = "camelCase")]
185struct TokenResponse {
186    access_token: AccessToken,
187    expire_time: DateTime<Utc>,
188}
189
190#[derive(Debug, Deserialize, Serialize)]
191pub struct StoredSecret {
192    access_token: AccessToken,
193    scopes: Scopes,
194    expire_time: DateTime<Utc>,
195}
196
197pub fn get_access_token(
198    gcloud_config: &GcloudConfig,
199    service_account: &Email,
200    lifetime: &Lifetime,
201    scopes: &Scopes,
202) -> anyhow::Result<AccessToken> {
203    let stored_secret = get_token_from_keyring(service_account);
204    match stored_secret {
205        Ok(s) => {
206            if &s.scopes != scopes {
207                println!("Scopes are not equal, getting a new token!");
208                let new_token =
209                    get_token_from_gcloud(service_account, lifetime, scopes, gcloud_config)?;
210                save_token_to_keyring(service_account, &new_token)?;
211                return Ok(new_token.access_token);
212            }
213
214            if s.expire_time <= Utc::now() {
215                println!("Token has expired, getting a new one!");
216                let new_token =
217                    get_token_from_gcloud(service_account, lifetime, scopes, gcloud_config)?;
218                save_token_to_keyring(service_account, &new_token)?;
219                return Ok(new_token.access_token);
220            }
221            return Ok(s.access_token);
222        }
223        Err(e) => match e {
224            keyring::Error::NoEntry => {
225                let new_token =
226                    get_token_from_gcloud(service_account, lifetime, scopes, gcloud_config)?;
227                save_token_to_keyring(service_account, &new_token)?;
228                return Ok(new_token.access_token);
229            }
230            other_error => panic!("failed to get access token: {:?}", other_error),
231        },
232    }
233}
234
235fn get_token_from_gcloud(
236    service_account: &Email,
237    lifetime: &Lifetime,
238    scopes: &Scopes,
239    gcloud_config: &GcloudConfig,
240) -> anyhow::Result<StoredSecret> {
241    let client: Client = Client::builder()
242        .user_agent(USER_AGENT)
243        .timeout(Duration::from_secs(15))
244        .build()?;
245
246    let url = format!(
247        "{}/projects/-/serviceAccounts/{}:generateAccessToken",
248        IAM_API, service_account
249    );
250
251    let mut headers = HeaderMap::new();
252    headers.insert(reqwest::header::ACCEPT, "application/json".parse()?);
253
254    let token_request = TokenRequest {
255        lifetime: format!("{}", lifetime),
256        scope: scopes.clone(),
257    };
258
259    let request = client
260        .post(url)
261        .bearer_auth(gcloud_config.access_token.as_ref())
262        .headers(headers)
263        .json(&token_request);
264
265    let response: TokenResponse = request.send()?.json()?;
266
267    Ok(StoredSecret {
268        access_token: response.access_token,
269        scopes: scopes.clone(),
270        expire_time: response.expire_time,
271    })
272}
273
274fn get_token_from_keyring(service_account: &Email) -> Result<StoredSecret, keyring::Error> {
275    let entry = Entry::new(env!("CARGO_PKG_NAME"), &service_account.0)?;
276    match entry.get_password() {
277        Ok(s) => {
278            let stored_secret: StoredSecret =
279                serde_json::from_str(&s).expect("failed to parse json from keyring");
280            Ok(stored_secret)
281        }
282        Err(e) => Err(e),
283    }
284}
285
286// fn delete_token_from_keyring(service_account: &Email) -> anyhow::Result<AccessToken> {
287//     todo!()
288// }
289
290fn save_token_to_keyring(
291    service_account: &Email,
292    stored_secret: &StoredSecret,
293) -> anyhow::Result<()> {
294    println!("Saving token to OS keyring!");
295    let secret_entry = serde_json::to_string(stored_secret)?;
296    let entry = Entry::new(env!("CARGO_PKG_NAME"), &service_account.0)?;
297    match entry.set_password(&secret_entry) {
298        Ok(_) => Ok(()),
299        Err(e) => Err(e.into()),
300    }
301}
302
303// TODO: support delegate chains? https://cloud.google.com/iam/docs/reference/credentials/rest/v1/projects.serviceAccounts/generateAccessToken