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
126impl 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
286fn 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