gemini_chat_api/
utils.rs

1//! Utility functions for cookie loading and file upload.
2
3use crate::enums::{upload_headers, Endpoint};
4use crate::error::{Error, Result};
5use reqwest::Client;
6use serde::Deserialize;
7use std::collections::HashMap;
8use std::path::Path;
9
10/// Cookie entry from browser export JSON format.
11#[derive(Debug, Deserialize)]
12struct CookieEntry {
13    name: String,
14    value: String,
15}
16
17/// Loads authentication cookies from a JSON file.
18///
19/// The file should be in the browser cookie export format:
20/// ```json
21/// [
22///   { "name": "__Secure-1PSID", "value": "..." },
23///   { "name": "__Secure-1PSIDTS", "value": "..." }
24/// ]
25/// ```
26///
27/// # Arguments
28/// * `cookie_path` - Path to the JSON cookie file
29///
30/// # Returns
31/// A tuple of (secure_1psid, secure_1psidts) values
32///
33/// # Errors
34/// Returns an error if the file is not found, invalid JSON, or missing required cookies.
35pub fn load_cookies(cookie_path: &str) -> Result<(String, String)> {
36    let path = Path::new(cookie_path);
37    if !path.exists() {
38        return Err(Error::Cookie(format!(
39            "Cookie file not found at path: {}",
40            cookie_path
41        )));
42    }
43
44    let content = std::fs::read_to_string(path)?;
45    let cookies: Vec<CookieEntry> = serde_json::from_str(&content)
46        .map_err(|e| Error::Cookie(format!("Invalid JSON format in cookie file: {}", e)))?;
47
48    let mut secure_1psid: Option<String> = None;
49    let mut secure_1psidts: Option<String> = None;
50
51    for cookie in cookies {
52        match cookie.name.to_uppercase().as_str() {
53            "__SECURE-1PSID" => secure_1psid = Some(cookie.value),
54            "__SECURE-1PSIDTS" => secure_1psidts = Some(cookie.value),
55            _ => {}
56        }
57    }
58
59    match (secure_1psid, secure_1psidts) {
60        (Some(psid), Some(psidts)) => Ok((psid, psidts)),
61        (None, _) => Err(Error::Cookie(
62            "Required cookie __Secure-1PSID not found".to_string(),
63        )),
64        (_, None) => Err(Error::Cookie(
65            "Required cookie __Secure-1PSIDTS not found".to_string(),
66        )),
67    }
68}
69
70/// Uploads a file to Google's Gemini server and returns its identifier.
71///
72/// # Arguments
73/// * `file_data` - The file content as bytes
74/// * `proxy` - Optional proxy URL
75///
76/// # Returns
77/// The file identifier string from the server
78///
79/// # Errors
80/// Returns an error if the upload fails.
81pub async fn upload_file(file_data: &[u8], proxy: Option<&str>) -> Result<String> {
82    let mut builder = Client::builder();
83
84    if let Some(proxy_url) = proxy {
85        builder = builder
86            .proxy(reqwest::Proxy::all(proxy_url).map_err(|e| Error::Upload(e.to_string()))?);
87    }
88
89    let client = builder.build().map_err(|e| Error::Upload(e.to_string()))?;
90
91    // Create multipart form with the file
92    let part = reqwest::multipart::Part::bytes(file_data.to_vec()).file_name("file");
93    let form = reqwest::multipart::Form::new().part("file", part);
94
95    let response: reqwest::Response = client
96        .post(Endpoint::Upload.url())
97        .headers(upload_headers())
98        .multipart(form)
99        .send()
100        .await
101        .map_err(|e| Error::Upload(e.to_string()))?;
102
103    if !response.status().is_success() {
104        return Err(Error::Upload(format!(
105            "Upload failed with status: {}",
106            response.status()
107        )));
108    }
109
110    let text = response
111        .text()
112        .await
113        .map_err(|e| Error::Upload(e.to_string()))?;
114    Ok(text)
115}
116
117/// Loads cookies from file and returns them as a HashMap for reqwest.
118pub fn cookies_to_map(secure_1psid: &str, secure_1psidts: &str) -> HashMap<String, String> {
119    let mut map = HashMap::new();
120    map.insert("__Secure-1PSID".to_string(), secure_1psid.to_string());
121    map.insert("__Secure-1PSIDTS".to_string(), secure_1psidts.to_string());
122    map
123}