diff options
| -rw-r--r-- | src/main.rs | 131 |
1 files changed, 94 insertions, 37 deletions
diff --git a/src/main.rs b/src/main.rs index fefa4cf..38a7ecb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ use base64::prelude::*; use yaml_rust2::parser::Parser; use yaml_rust2::Event; +use std::collections::HashMap; use std::ops::Range; fn help(cmd: &str) { @@ -9,32 +10,73 @@ fn help(cmd: &str) { Search FILE for docker images and suggest updates Options: ---auth <USER> <PASS> Authenticate against registry", +--auth <REGISTRY> <USER> <PASS> Authenticate against REGISTRY (repeatable)", cmd ); } +struct AuthInfo { + host: String, + username: String, + password: String, +} + +enum AuthStage { + Unauthorized, + Authorized(String), +} + +struct AuthState { + info: AuthInfo, + stage: AuthStage, +} + struct Auth { - header: Option<String>, + states: HashMap<String, AuthState> } impl Auth { - fn new() -> Self { - return Auth{ - header: None, - }; + fn new(infos: Vec<AuthInfo>) -> Self { + let mut states = HashMap::new(); + for info in infos { + states.insert(info.host.clone(), AuthState { + info: info, + stage: AuthStage::Unauthorized, + }); + } + + return Auth { + states + } } - fn parse(&mut self, username: &str, password: &str) { - self.header = Some(format!("Basic {}", BASE64_STANDARD.encode(format!("{}:{}", username, password)))); + fn first(&mut self, host: &str) -> Option<String> { + if let Some(state) = self.states.get_mut(host) { + match state.stage { + AuthStage::Unauthorized => { return None }, + AuthStage::Authorized(ref x) => { return Some(x.clone()); }, + } + } + + return None; } - fn apply<B>(&self, req: ureq::RequestBuilder<B>) -> ureq::RequestBuilder<B> { - if let Some(header) = &self.header { - return req.header("Authorization", header.clone()); - } else { - return req; + fn authenticate<T>(&mut self, host: &str, previous_response: &ureq::http::Response<T>) -> Option<String> { + if let Some(state) = self.states.get_mut(host) { + match previous_response.headers().get("WWW-Authenticate").map(|x| x.to_str()) { + Some(Ok("basic")) => { + let header = format!("Basic {}", BASE64_STANDARD.encode(format!("{}:{}", state.info.username, state.info.password))); + state.stage = AuthStage::Authorized(header.clone()); + + return Some(header); + }, + Some(Ok("bearer")) => todo!(), + Some(_) => return None, + None => return None, + } } + + return None; } } @@ -153,7 +195,7 @@ impl DockerRef{ let mut tag = None; if let Some(idx) = file[string_range.clone()].rfind(":") { tag = Some(string_range.start+idx+1..string_range.end); - string_range.end = idx; + string_range.end = string_range.start+idx; } let mut registry = None; @@ -188,37 +230,45 @@ struct Update { content: String, } -fn fetch_new_image(file: &str, auth: &Auth, img: DockerRef, edits: &mut Vec<Update>) { - let mut resulting_ref = img.clone(); - - dbg!(&img); - let registry = img.registry.map(|x| &file[x]).unwrap_or("https://registry.jnsn.dev/"); +fn fetch_new_image(file: &str, auth: &mut Auth, img: DockerRef, edits: &mut Vec<Update>) { + let registry = img.registry.map(|x| &file[x]).unwrap_or("registry.jnsn.dev/"); let tag = img.tag.map(|x| &file[x]).unwrap_or("latest"); - if tag == "latest" || img.digest.is_some() { + // Find the digest for the newest image + if let Some(digest) = img.digest { let url = format!("https://{}/v2/{}/manifests/{}", registry, &file[img.image], tag); dbg!(&url); - let mut response = auth.apply(ureq::get(url)) + let mut response = ureq::get(&url) + .header("Authorization", auth.first(registry)) .call().unwrap(); + if response.status() == 401 { + response = ureq::get(&url) + .header("Authorize", auth.authenticate(registry, &response).unwrap()) + .call().unwrap(); + } + let body: tinyjson::JsonValue = response .body_mut() - .read_to_string().unwrap().parse().unwrap(); + .read_to_string() + .unwrap() + .parse() + .unwrap(); let media_type : &String = body["mediaType"].get().unwrap(); assert!(media_type == "application/vnd.docker.distribution.manifest.v2+json"); - let digest = &response.headers()["docker-content-digest"]; - // resulting_ref.digest = Some(digest.to_str().unwrap().to_string()); - - // dbg!(&resulting_ref.raw[resulting_ref.digest.as_ref().unwrap().clone()]); - dbg!(&resulting_ref); + let image_ref = { + let prefix = &file[img.full_range.start..digest.start]; + let digest = &response.headers()["docker-content-digest"].to_str().unwrap(); + format!("{}{}", prefix, digest) + }; edits.push(Update{ position: img.full_range, - content: "AHH".to_string(), + content: image_ref, }); } } @@ -228,17 +278,21 @@ fn main() { let mut it = argv.iter(); let cmd = &it.next().unwrap(); - let mut auth = Auth::new(); + let mut auths = vec![]; let mut positional: Vec<&str> = vec!(); loop { match it.next().map(|x| x.as_str()) { None => break, Some("--auth") => { - if let Some(username) = it.next() && let Some(password) = it.next() { - auth.parse(username, password); + if let Some(registry) = it.next() && let Some(username) = it.next() && let Some(password) = it.next() { + auths.push(AuthInfo { + host: registry.clone(), + username: username.clone(), + password: password.clone(), + }); } else { - println!("Error: --auth requires two parameters"); + println!("Error: --auth requires three parameters"); help(cmd); std::process::exit(1); } @@ -251,11 +305,14 @@ fn main() { }; } + if positional.len() != 1 { panic!("Bad arguments"); } let file = positional[0]; + let mut auth = Auth::new(auths); + let file_content = &std::fs::read_to_string(file).unwrap(); let yaml = Parser::new_from_str(&file_content); let images : Vec<_> = scan_yaml_for_images(yaml); @@ -264,16 +321,16 @@ fn main() { let images : Vec<_> = images.iter() .map(|x| DockerRef::parse(&file_content, x)) - .map(|x| fetch_new_image(&file_content, &auth, x, &mut edits)) + .map(|x| fetch_new_image(&file_content, &mut auth, x, &mut edits)) .collect(); dbg!(&images); dbg!(&edits); // dbg!(&file_content[images[0].digest.as_ref().unwrap().clone()]); - let body: String = auth.apply(ureq::get("https://registry.jnsn.dev/v2/autobrr/tags/list")) - .call().unwrap() - .body_mut() - .read_to_string().unwrap(); + // let body: String = auth.apply(ureq::get("https://registry.jnsn.dev/v2/autobrr/tags/list")) + // .call().unwrap() + // .body_mut() + // .read_to_string().unwrap(); // dbg!(body); } |