path: root/src/main.rs

mod parser;
use crate::parser::*;

use base64::prelude::*;
use rand::distr::{Alphanumeric, SampleString};
use yaml_rust2::parser::Parser;
use yaml_rust2::Event;
use std::collections::HashMap;
use std::ops::Range;
use std::sync;
use std::io::Read;
use std::io::Seek;
use std::io::Write;

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

#[derive(Debug)]
enum AuthMethod {
    Basic,
    Bearer{realm: String, service: String, scope: String},
}

impl AuthMethod {
    fn from_header(header: &ureq::http::HeaderValue) -> Option<Self> {
        let header_str = header.to_str().unwrap();
        let parse = parse_authenticate_header(header_str).unwrap();

        if &header_str[parse.scheme.clone()] == "Basic" {
            return Some(AuthMethod::Basic);
        }

        if &header_str[parse.scheme.clone()] == "Bearer" {
            return Some(AuthMethod::Bearer{
                realm: header_str[parse.realm.unwrap()].to_string(),
                scope: header_str[parse.scope.unwrap()].to_string(),
                service: header_str[parse.service.unwrap()].to_string(),
            });
        }

        return None;
    }
}

struct AuthInfo {
    host: String,
    username: String,
    password: String,
}

enum AuthStage {
    Idle,
    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::Idle,
            });
        }

        return Auth {
            states
        }
    }
}

impl AuthState{
    fn add_to_request<T>(&self, req: ureq::RequestBuilder<T>) -> ureq::RequestBuilder<T> {
        if let AuthStage::Authorized(x) = &self.stage {
            return req.header("Authorization", x);
        }

        return req
    }

    fn authenticate(&mut self, response: &ureq::http::Response<ureq::Body>) -> Result<(), String> {
        match self.stage {
            AuthStage::Authorized(_) => 
                // The token must have expired
                self.stage = AuthStage::Idle,
            AuthStage::Idle => {},
        }

        let auth_header = response.headers().get("www-authenticate").unwrap();
        match AuthMethod::from_header(auth_header) {
            Some(AuthMethod::Basic) => {
                let basic_auth = format!("Basic {}", BASE64_STANDARD.encode(format!("{}:{}", self.info.username, self.info.password)));
                self.stage = AuthStage::Authorized(basic_auth);

                return Ok(());
            },
            Some(AuthMethod::Bearer{realm, scope, service}) => {
                let basic_auth = format!("Basic {}", BASE64_STANDARD.encode(format!("{}:{}", self.info.username, self.info.password)));

                let url = format!("{}?service={}&scope={}", realm, service, scope);

                let body = ureq::get(url)
                    .config().http_status_as_error(false).build()
                    .header("Authorization", basic_auth)
                    .call();

                let body = body.map_err(|x| format!("Server {} authentication request failed: {}", self.info.host, x.to_string()))?
                    .body_mut().read_to_string().expect(format!("Server {} responded with something non-string like", self.info.host).as_str());

                let body = body.parse::<tinyjson::JsonValue>()
                    .unwrap();

                let token = body["token"].get::<String>().unwrap();
                self.stage = AuthStage::Authorized(format!("Bearer {}", token));
                return Ok(());

            },
            None => Err(format!("Server {} provided us with a challenge, but we didn't understand it", self.info.host)),
        }
    }
}

fn perform_registry_request(registry: &str, url: &str, accept: &'static str, auth: &mut Auth) -> Result<ureq::http::Response<ureq::Body>, String> {
    let mut state = auth.states.get_mut(registry);

    let url = format!("https://{}{}", registry, &url);

    for _ in 0..2 {
        let mut request = ureq::get(&url)
            .header("Accept", accept)
            .config().http_status_as_error(false).build();

        if let Some(ref state) = state {
            request = state.add_to_request(request);
        }

        let response = request.call().unwrap();

        if response.status() == 401 {
            if let Some(ref mut state) = state {
                state.authenticate(&response)?;
                continue;
            } else {
                return Err(format!("Server {} returned 401 but we have no credentials", registry));
            }
        } else if response.status() != 200 {
            return Err(format!("Unexpected status code: {}", response.status()))
        }

        return Ok(response);
    }

    return Err("Authorization failed".to_string());
}

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

#[derive(Debug)]
struct ManifestFile {
    image_tags: Vec<Range<usize>>,
}

impl ManifestFile {
    fn parse(content: &str) -> Self {
        let mut yaml = Parser::new_from_str(content);

        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 {
                                images.push(mark.index()..mark.index() + txt.len());
                            }
                            scope.pop();
                        },
                        // This should only happen for entirely empty documents
                        YContext::InDocument => assert!(txt == ""),

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

        return Self{
            image_tags: images,
        };
    }
}

#[derive(Debug, Clone)]
struct DockerRef {
    registry: Option<Range<usize>>,
    image: Range<usize>,
    tag: Option<Range<usize>>,
    digest: Option<Range<usize>>,
}

impl DockerRef{
    fn parse(content: &str, chunk: &Range<usize>) -> DockerRef {
        let mut string_range = chunk.clone();

        let mut digest = None;
        if let Some(idx) = content[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) = content[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) = content[string_range.clone()].find("/") {
            let head = &content[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 {
            registry,
            image,
            tag,
            digest,
        };
    }
}

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

#[derive(Debug)]
enum VersionPart {
    String(String),
    Number(u64),
    Hash,
}

#[derive(Debug)]
struct VersionPattern {
    parts: Vec<VersionPart>,
}

#[derive(Debug)]
enum CompareOutcome {
    Higher,
    Lower,

    Incompatible,
    Identical,
}

impl VersionPattern {
    fn parse(tag: &str) -> Self {
        static RE: sync::LazyLock<regex::Regex> = sync::LazyLock::new(|| regex::Regex::new(r"(?<hash>[a-f0-9]{32})|(?<str>[^0-9]+)|(?<num>[0-9]+)").unwrap());
        let mut parts = vec![];
        for it in RE.captures_iter(tag) {
            if let Some(x) = it.name("str") {
                parts.push(VersionPart::String(x.as_str().to_string()));
            } else if let Some(x) = it.name("num") {
                parts.push(VersionPart::Number(x.as_str().parse().unwrap()));
            } else if let Some(_) = it.name("hash") {
                parts.push(VersionPart::Hash);
            }
        }

        return VersionPattern {
            parts,
        }
    }

    fn compare(&self, other: &Self) -> CompareOutcome {
        if self.parts.len() != other.parts.len() {
            return CompareOutcome::Incompatible;
        }

        for (self_part, other_part) in self.parts.iter().zip(other.parts.iter()) {
            match (self_part, other_part) {
                (VersionPart::String(x1), VersionPart::String(x2)) => if x1 != x2 { return CompareOutcome::Incompatible },
                (VersionPart::String(_), _) => return CompareOutcome::Incompatible,
                (VersionPart::Number(x1), VersionPart::Number(x2)) => {
                    if x1 > x2 {
                        return CompareOutcome::Lower
                    } else if x1 > x2 {
                        return CompareOutcome::Higher
                    }
                },
                (VersionPart::Number(_), _) => return CompareOutcome::Incompatible,
                (VersionPart::Hash, VersionPart::Hash) => {},
                (VersionPart::Hash, _) => return CompareOutcome::Incompatible,
            }
        }

        return CompareOutcome::Identical;
    }
}

fn update_images(file: &str, auth: &mut Auth, img: DockerRef, edits: &mut Vec<FilePatch>) -> Result<(), String> {
    let registry = img.registry.map(|x| &file[x]).unwrap_or("registry.hub.docker.com");
    let mut tag = img.tag.as_ref().map(|x| file[x.clone()].to_string());

    if let Some(ref tag_str) = tag {
        let mut current = VersionPattern::parse(&tag_str);

        let mut new_tag = None;

        let mut url = format!("/v2/{}/tags/list", &file[img.image.clone()]);
        loop {
            let mut response = perform_registry_request(registry, &url, "application/vnd.oci.image.index.v1+json", auth).unwrap();

            let content_type = response.headers()["Content-Type"].to_str().unwrap();
            if !content_type.starts_with("application/json") {
                return Err(format!("Unexpected Content-Type: {}", content_type));
            }

            let body = response.body_mut().read_to_string().unwrap();
            let body = body.parse::<tinyjson::JsonValue>().unwrap();

            for it in body["tags"].get::<Vec<tinyjson::JsonValue>>().unwrap().iter() {
                let candidate_str = it.get::<String>().unwrap();
                let candidate = VersionPattern::parse(candidate_str);
                match current.compare(&candidate) {
                    CompareOutcome::Higher => {
                        new_tag = Some(candidate_str.clone());
                        current = candidate;
                    },
                    CompareOutcome::Lower => {},
                    CompareOutcome::Incompatible => {},
                    CompareOutcome::Identical => {},
                }
            }

            if let Some(link_header) = response.headers().get("link") {
                let link_str = link_header.to_str().unwrap();
                let link = extract_next_page(&link_str).unwrap();
                url = link_str[link.next_uri.unwrap().clone()].to_string();
            } else {
                break;
            }
        }

        if let Some(new_tag) = new_tag {
            tag = Some(new_tag.clone());
            edits.push(FilePatch {
                position: img.tag.unwrap(),
                content: new_tag,
            });
        }
    }

    if let Some(ref digest) = img.digest {
        // Find the digest for the selected tag
        let url = format!("/v2/{}/manifests/{}", &file[img.image.clone()], tag.as_deref().unwrap_or("latest"));
        let response = perform_registry_request(registry, &url, "application/vnd.oci.image.manifest.v1+json,application/vnd.oci.image.index.v1+json", auth).unwrap();

        let digest_string = response.headers()["docker-content-digest"].to_str().unwrap().to_string();

        if file[digest.clone()] != digest_string {
            edits.push(FilePatch{
                position: img.digest.unwrap(),
                content: digest_string,
            });
        }
    }

    return Ok(());
}

enum Output {
    Overlay(std::fs::File, std::path::PathBuf, std::path::PathBuf),
    Stdout(std::io::Stdout),
}

impl Output {
    fn overlay_file(infile_path: &std::path::Path) -> Self {
        loop {
            let mut filename = std::ffi::OsString::new();
            filename.push(infile_path.file_name().unwrap());
            filename.push(std::ffi::OsStr::new(".edit"));
            filename.push(Alphanumeric.sample_string(&mut rand::rng(), 16));
            let outfile_path = Some(infile_path.parent().unwrap().join(filename));
            if let Ok(output_file) = std::fs::File::create_new(outfile_path.as_ref().unwrap()) {
                return Output::Overlay(output_file, outfile_path.unwrap().into(), infile_path.into());
            }
        }
    }

    fn commit(&mut self) {
        match self {
            Output::Overlay(_, outfile_path, infile_path) => {
                std::fs::remove_file(&infile_path).unwrap();
                std::fs::rename(&outfile_path, &infile_path).unwrap();
            }
            Output::Stdout(_) => {}
        }
    }
}

impl std::ops::Deref for Output {
    type Target = dyn Write;

    fn deref(&self) -> &Self::Target {
        match self {
            Output::Overlay(file, _, _) => file,
            Output::Stdout(stdout) => stdout,
        }
    }
}

impl std::ops::DerefMut for Output {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            Output::Overlay(file, _, _) => file,
            Output::Stdout(stdout) => stdout,
        }
    }
}

impl std::io::Write for Output {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        return (**self).write(buf);
    }

    fn flush(&mut self) -> std::io::Result<()> {
        return (**self).flush();
    }

    fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
        return (**self).write_vectored(bufs);
    }

    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
        return (**self).write_all(buf);
    }

    fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
        return (**self).write_fmt(args);
    }
}

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!();
    let mut overwrite = false;

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


    if positional.len() != 1 {
        panic!("Bad arguments");
    }
    let infile_path = std::path::PathBuf::from(positional[0]);
    let mut auth = Auth::new(auths);

    let inpaths = if infile_path.is_dir() {
        let mut unsearched = vec![infile_path.clone()];
        let mut paths = vec![];

        while let Some(next) = unsearched.pop() {
            for child in next.read_dir().unwrap() {
                let child = child.unwrap();
                let path = child.path();

                let ft = child.file_type().unwrap();
                if ft.is_dir() {
                    unsearched.push(path);
                    continue;
                }

                if let Some(ext) = path.extension() {
                    if ext == "yaml" {
                        paths.push(path);
                    }
                }
            }
        }

        paths
    } else {
        vec![infile_path]
    };

    for infile_path in inpaths {
        let mut file = std::fs::File::open(&infile_path).unwrap();
        let mut file_content = String::new();
        file.read_to_string(&mut file_content).unwrap();

        let images = ManifestFile::parse(&file_content);

        let mut failed = None;

        let mut edits = vec![];
        for ref image in images.image_tags {
            let image_ref = DockerRef::parse(&file_content, image);
            if let Err(msg) = update_images(&file_content, &mut auth, image_ref, &mut edits) {
                failed = Some(msg);
                break;
            }
        }

        if let Some(msg) = failed {
            println!("{}: {} ", infile_path.display(), msg);
            continue;
        }

        let mut out = if overwrite {
            Output::overlay_file(&infile_path)
        } else {
            Output::Stdout(std::io::stdout())
        };

        let mut current_position = 0;
        file.seek(std::io::SeekFrom::Start(0)).unwrap();
        let mut file_block = file.take(0);
        for edit in edits {
            if edit.position.start > current_position {
                file_block.set_limit((edit.position.start - current_position) as u64);
                std::io::copy(&mut file_block, &mut out).unwrap();
            }

            out.write_all(edit.content.as_bytes()).unwrap();

            // We should have been able to just seek to the position.end here, but that doesn't work
            // for whatever reason. What we can do is calculate the amount to skip ahead, and then do
            // that.
            let skip = (edit.position.end - edit.position.start) as i64;
            file_block.set_limit(skip as u64);
            file_block.seek(std::io::SeekFrom::Current(skip)).unwrap();
            current_position = edit.position.end;
        }

        file_block.set_limit(u64::MAX);
        std::io::copy(&mut file_block, &mut out).unwrap();

        out.commit();
    }
}