1use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
2use serde::{Deserialize, Serialize};
3
4#[path = "google/protos.rs"]
5pub mod google_cloud_ai;
6
7#[derive(Debug, Serialize, Deserialize)]
8struct Token {
9 access_token: String,
10}
11
12#[derive(Debug, Serialize, Deserialize)]
13struct Claims {
14 iss: String,
15 scope: String,
16 aud: String,
17 exp: usize,
18 iat: usize,
19}
20
21pub async fn get_token() -> String {
22 let client = reqwest::Client::new();
23
24 let response = client.get("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token")
25 .header("Metadata-Flavor", "Google")
26 .send()
27 .await
28 .unwrap();
29
30 let result = response.json::<Token>().await.unwrap();
31
32 return result.access_token;
33}
34
35pub async fn get_service_account_token(email: String, rsa_pem: String) -> String {
36 let timestamp = chrono::Utc::now();
37
38 let iat_value = timestamp.timestamp().try_into().unwrap();
39 let exp_value = (timestamp.timestamp() + 3600).try_into().unwrap();
40
41 let claims = Claims {
42 iss: email,
43 scope: String::from("https://www.googleapis.com/auth/datastore"),
44 aud: String::from("https://oauth2.googleapis.com/token"),
45 exp: exp_value,
46 iat: iat_value,
47 };
48
49 let mut headers = Header::new(Algorithm::RS256);
50
51 headers.typ = Some("JWT".to_string());
52
53 let token = encode(
54 &headers,
55 &claims,
56 &EncodingKey::from_rsa_pem(rsa_pem.as_ref()).unwrap(),
57 )
58 .unwrap();
59
60 let client = reqwest::Client::new();
61
62 let resp = client
63 .post("https://oauth2.googleapis.com/token")
64 .header("Content-Type", "application/x-www-form-urlencoded")
65 .body(format!(
66 "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion={}",
67 token
68 ))
69 .send()
70 .await
71 .unwrap();
72
73 let result = resp.json::<Token>().await.unwrap();
74
75 return result.access_token;
76}