trifid/tfclient/build.rs

use flate2::read::GzDecoder;
use reqwest::blocking::Response;
use reqwest::header::HeaderMap;
use std::fs;
use std::fs::{remove_file, File};
use std::io::{Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::path::Path;
use std::process::{Command, Output};
use tar::Archive;

#[derive(serde::Deserialize, Debug)]
struct GithubRelease {
    name: String,
    assets: Vec<GithubReleaseAsset>,
}

#[derive(serde::Deserialize, Debug)]
struct GithubUser {}

#[derive(serde::Deserialize, Debug)]
struct GithubReleaseAsset {
    browser_download_url: String,
    name: String,
    size: i64,
}

fn main() {
    if Path::new(&format!(
        "{}/{}",
        std::env::var("OUT_DIR").unwrap(),
        "noredownload"
    ))
    .exists()
        && std::env::var("TFBUILD_FORCE_REDOWNLOAD").is_err()
    {
        println!("noredownload exists and TFBUILD_FORCE_REDOWNLOAD is not set. Not redoing build process.");
        return;
    }

    println!("[*] Fetching nebula releaseinfo...");

    let mut headers = HeaderMap::new();
    let mut has_api_key = false;

    if let Ok(api_key) = std::env::var("GH_API_KEY") {
        headers.insert(
            "Authorization",
            format!("Bearer {}", api_key).parse().unwrap(),
        );
        has_api_key = true;
    }

    let client = reqwest::blocking::Client::builder()
        .user_agent("curl/7.57.1")
        .default_headers(headers)
        .build()
        .unwrap();

    let resp: Response = client
        .get("https://api.github.com/repos/slackhq/nebula/releases/latest")
        .send()
        .unwrap();

    if resp
        .headers()
        .get("X-Ratelimit-Remaining")
        .unwrap()
        .to_str()
        .unwrap()
        == "0"
    {
        println!("You've been ratelimited from the GitHub API. Wait a while (1 hour)");
        if !has_api_key {
            println!("You can also set a GitHub API key with the environment variable GH_API_KEY, which will increase your ratelimit ( a lot )");
        }
        panic!("Ratelimited");
    }

    let release: GithubRelease = resp.json().unwrap();

    println!("[*] Fetching target triplet...");
    let target = std::env::var("TARGET").unwrap();

    println!("[*] Compiling for target {}", target);

    let target_file = match target.as_str() {
        "x86_64-apple-darwin" | "aarch64-apple-darwin" => "nebula-darwin",
        "x86_64-unknown-freebsd" => "nebula-freebsd-amd64",
        "x86_64-unknown-linux-gnu" | "x86_64-unknown-linux-musl" => "nebula-linux-amd64",
        "armv5te-unknown-linux-gnueabi" => "nebula-linux-arm-5",
        "arm-unknown-linux-gnueabi" | "arm-unknown-linux-gnueabihf" => "nebula-linux-arm-6",
        "armv7-unknown-linux-gnueabihf" | "armv7-unknown-linux-gnueabi" => "nebula-linux-arm-7",
        "aarch64-unknown-linux-gnu" | "aarch64-unknown-linux-musl" => "nebula-linux-arm64",
        "x86_64-pc-windows-msvc" => "nebula-windows-amd64",
        "aarch64-pc-windows-msvc" => "nebula-windows-arm64",
        _ => {
            println!("This architecture is not supported yet :(");
            println!("Nebula has a limited set of architectures it is able to function on.");
            println!("tfclient can only be compiled on these architectures.");
            println!("See https://github.com/slackhq/nebula/releases for a list of supported architectures");
            println!("Is your system supported by Nebula? Shoot a message to the mailing list. Include the target triplet (above) in your response, as well as a link to the functioning Nebula binary. We will happily add your machine to the list!");
            panic!("Unsupported architecture");
        }
    };

    println!("[*] Embedding {} {}", target_file, release.name);

    let download = release
        .assets
        .iter()
        .find(|r| r.name == format!("{}.tar.gz", target_file))
        .expect("That architecture isn't avaliable :(");

    println!(
        "[*] Downloading {}.tar.gz ({}, {} bytes) from {}",
        target_file, target, download.size, download.browser_download_url
    );

    let response = reqwest::blocking::get(&download.browser_download_url).unwrap();
    let content = response.bytes().unwrap().to_vec();
    let bytes = content.as_slice();

    let tar = GzDecoder::new(bytes);
    let mut archive = Archive::new(tar);
    let entries = archive.entries().unwrap();

    let mut nebula_bin = vec![];
    let mut nebula_cert_bin = vec![];
    let mut shasum = vec![];

    for entry in entries {
        let mut entry = entry.unwrap();
        if entry.path().unwrap() == Path::new("nebula")
            || entry.path().unwrap() == Path::new("nebula.exe")
        {
            nebula_bin.reserve(entry.size() as usize);
            entry.read_to_end(&mut nebula_bin).unwrap();
        } else if entry.path().unwrap() == Path::new("nebula-cert")
            || entry.path().unwrap() == Path::new("nebula-cert.exe")
        {
            nebula_cert_bin.reserve(entry.size() as usize);
            entry.read_to_end(&mut nebula_cert_bin).unwrap();
        } else if entry.path().unwrap() == Path::new("SHASUM256.txt") {
            shasum.reserve(entry.size() as usize);
            entry.read_to_end(&mut shasum).unwrap();
        }
    }

    if nebula_bin.is_empty() {
        panic!("[x] Release did not contain nebula binary");
    }
    if nebula_cert_bin.is_empty() {
        panic!("[x] Release did not contain nebula_cert binary");
    }

    let mut nebula_file =
        File::create(format!("{}/nebula.bin", std::env::var("OUT_DIR").unwrap())).unwrap();
    nebula_file.write_all(&nebula_bin).unwrap();

    codegen_version(&nebula_bin, "nebula.bin", "NEBULA");

    let mut nebula_cert_file = File::create(format!(
        "{}/nebula_cert.bin",
        std::env::var("OUT_DIR").unwrap()
    ))
    .unwrap();
    nebula_cert_file.write_all(&nebula_cert_bin).unwrap();

    codegen_version(&nebula_cert_bin, "nebula_cert.bin", "NEBULA_CERT");

    // Indicate to cargo and ourselves that we have already downloaded and codegenned
    File::create(format!(
        "{}/{}",
        std::env::var("OUT_DIR").unwrap(),
        "noredownload"
    ))
    .unwrap();
    println!("cargo:rerun-if-changed=build.rs");
}

fn codegen_version(bin: &[u8], fp: &str, name: &str) {
    // get version
    let output = execim(bin, &vec!["-version"]);
    let stdout = output.stdout;
    let stdout_str = String::from_utf8(stdout).unwrap();
    if !stdout_str.starts_with("Version: ") {
        panic!("Binary did not have expected version output. Unable to get version info.");
    }
    let mut version = stdout_str.split(' ').collect::<Vec<&str>>()[1].to_string();
    version.pop();

    let code = format!("// This code was automatically @generated by build.rs. It should not be modified.\npub const {}_BIN: &[u8] = include_bytes!(concat!(env!(\"OUT_DIR\"), \"/{}\"));\npub const {}_VERSION: &str = \"{}\";", name, fp, name, version);

    let mut file =
        File::create(format!("{}/{}.rs", std::env::var("OUT_DIR").unwrap(), fp)).unwrap();
    file.write_all(code.as_bytes()).unwrap();
}

#[cfg(not(unix))]
fn execim(buf: &[u8], args: &Vec<&str>) -> Output {
    let mut file = File::create("tmpexec.bin").unwrap();
    file.write_all(buf).unwrap();
    std::mem::drop(file);
    let output = Command::new("./tmpexec.bin").args(args).output().unwrap();
    remove_file("./tmpexec.bin").unwrap();
    output
}

#[cfg(unix)]
fn execim(buf: &[u8], args: &Vec<&str>) -> Output {
    let mut file = File::create("tmpexec.bin").unwrap();
    file.write_all(buf).unwrap();
    let metadata = file.metadata().unwrap();
    let mut permissions = metadata.permissions();
    permissions.set_mode(0o0755);
    fs::set_permissions("./tmpexec.bin", permissions).unwrap();
    std::mem::drop(file);
    let output = Command::new("./tmpexec.bin").args(args).output().unwrap();
    remove_file("./tmpexec.bin").unwrap();
    output
}