path: root/src/main.rs

use base64::prelude::*;
use yaml_rust2::parser::Parser;
use yaml_rust2::Event;
use std::collections::HashMap;
use std::ops::Range;

fn help(cmd: &str) {
    println!(
        "{} [options] [--] <FILE>
Search FILE for docker images and suggest updates

Options:
--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 {
    states: HashMap<String, AuthState>
}

impl Auth {
    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 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 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;
    }
}

enum YContext {
    InDocument,
    InObject,
    InSequence,
    InValue(bool),
}

#[derive(Debug)]
struct Chunk {
    position: Range<usize>,
}

fn scan_yaml_for_images<T: Iterator<Item = char>>(mut yaml: Parser<T>) -> Vec<Chunk> {
    let mut images = vec!();

    let mut scope = vec!();

    loop {
        let (ev, mark) = yaml.next_token().unwrap();
        match ev {
            Event::StreamStart => {}
            Event::StreamEnd => { break; }
            Event::DocumentStart => { scope.push(YContext::InDocument); }
            Event::DocumentEnd => {
                assert!(matches!(scope.pop().unwrap(), YContext::InDocument));
                scope.pop_if(|x| matches!(x, YContext::InValue(_)));
            },
            Event::MappingStart(_, _) => { scope.push(YContext::InObject); },
            Event::MappingEnd => {
                assert!(matches!(scope.pop().unwrap(), YContext::InObject));
                scope.pop_if(|x| matches!(x, YContext::InValue(_)));
            },
            Event::SequenceStart(_, _) => { scope.push(YContext::InSequence); },
            Event::SequenceEnd => { 
                assert!(matches!(scope.pop().unwrap(), YContext::InSequence));
                scope.pop_if(|x| matches!(x, YContext::InValue(_)));
            },

            Event::Scalar(ref txt, _, _, _) => {
                let parent = scope.last().unwrap();
                match parent {
                    YContext::InObject => {
                        // We are the key of a mapping, which means the next even is the value
                        scope.push(YContext::InValue(txt == "image"));
                    },
                    YContext::InSequence => {},
                    YContext::InValue(img) => {
                        if *img {
                            let next_idx = images.len();
                            images.push(Chunk{
                                position: mark.index()..mark.index() + txt.len(),
                            });
                        }
                        scope.pop();
                    },

                    _ => panic!(),
                }
            },
            x => todo!("{:?}", x),
        }
    }

    return images;
}

pub trait SubsliceOffset<T> {
    fn subslice_range(&self, inner: &Self) -> Option<std::ops::Range<usize>>;
}

impl<T> SubsliceOffset<T> for [T] {
    fn subslice_range(&self, subslice: &[T]) -> Option<std::ops::Range<usize>> {
        if size_of::<T>() == 0 {
            panic!("elements are zero-sized");
        }

        let self_start = self.as_ptr().addr();
        let subslice_start = subslice.as_ptr().addr();

        let byte_start = subslice_start.wrapping_sub(self_start);

        if !byte_start.is_multiple_of(size_of::<T>()) {
            return None;
        }

        let start = byte_start / size_of::<T>();
        let end = start.wrapping_add(subslice.len());

        if start <= self.len() && end <= self.len() { Some(start..end) } else { None }
    }
}

#[derive(Debug, Clone)]
struct DockerRef {
    full_range: Range<usize>,

    registry: Option<Range<usize>>,
    image: Range<usize>,
    tag: Option<Range<usize>>,
    digest: Option<Range<usize>>,
}

impl DockerRef{
    fn parse(file: &str, chunk: &Chunk) -> DockerRef {
        let mut string_range = chunk.position.clone();

        let mut digest = None;
        if let Some(idx) = file[string_range.clone()].rfind("@") {
            digest = Some(string_range.start+idx+1..string_range.end);
            string_range.end = string_range.start+idx;
        }

        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 = string_range.start+idx;
        }

        let mut registry = None;
        let image;
        if let Some(idx) = file[string_range.clone()].find("/") {
            let head = &file[string_range.clone()][..idx];
            if head.contains(":") || head.contains(".") {
                registry = Some(string_range.start..string_range.start+idx);
                image = string_range.start+idx+1..string_range.end;
            } else {
                registry = None;
                image = string_range;
            }
        } else {
            image = string_range;
        }

        return DockerRef {
            full_range: chunk.position.clone(),

            registry,
            image,
            tag,
            digest,
        };
    }
}

#[derive(Debug)]
struct Update {
    position: Range<usize>,
    content: String,
}

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");

    // 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 = 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();


        let media_type : &String = body["mediaType"].get().unwrap();
        assert!(media_type == "application/vnd.docker.distribution.manifest.v2+json");

        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: image_ref,
        });
    }
}

fn main() {
    let argv: Vec<String> = std::env::args().collect();
    let mut it = argv.iter();
    let cmd = &it.next().unwrap();

    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(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 three parameters");
                    help(cmd);
                    std::process::exit(1);
                }
            },
            Some("-h") | Some("--help") => {
                help(cmd);
                std::process::exit(0);
            },
            Some(arg) => positional.push(arg),
        };
    }


    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);

    let mut edits = vec![];

    let images : Vec<_> = images.iter()
        .map(|x| DockerRef::parse(&file_content, x))
        .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();
    // dbg!(body);
}