trifid/tfclient/build.rs

/*
[
  {
    "url": "https://api.github.com/repos/octocat/Hello-World/releases/1",
    "html_url": "https://github.com/octocat/Hello-World/releases/v1.0.0",
    "assets_url": "https://api.github.com/repos/octocat/Hello-World/releases/1/assets",
    "upload_url": "https://uploads.github.com/repos/octocat/Hello-World/releases/1/assets{?name,label}",
    "tarball_url": "https://api.github.com/repos/octocat/Hello-World/tarball/v1.0.0",
    "zipball_url": "https://api.github.com/repos/octocat/Hello-World/zipball/v1.0.0",
    "id": 1,
    "node_id": "MDc6UmVsZWFzZTE=",
    "tag_name": "v1.0.0",
    "target_commitish": "master",
    "name": "v1.0.0",
    "body": "Description of the release",
    "draft": false,
    "prerelease": false,
    "created_at": "2013-02-27T19:35:32Z",
    "published_at": "2013-02-27T19:35:32Z",
    "author": {
      "login": "octocat",
      "id": 1,
      "node_id": "MDQ6VXNlcjE=",
      "avatar_url": "https://github.com/images/error/octocat_happy.gif",
      "gravatar_id": "",
      "url": "https://api.github.com/users/octocat",
      "html_url": "https://github.com/octocat",
      "followers_url": "https://api.github.com/users/octocat/followers",
      "following_url": "https://api.github.com/users/octocat/following{/other_user}",
      "gists_url": "https://api.github.com/users/octocat/gists{/gist_id}",
      "starred_url": "https://api.github.com/users/octocat/starred{/owner}{/repo}",
      "subscriptions_url": "https://api.github.com/users/octocat/subscriptions",
      "organizations_url": "https://api.github.com/users/octocat/orgs",
      "repos_url": "https://api.github.com/users/octocat/repos",
      "events_url": "https://api.github.com/users/octocat/events{/privacy}",
      "received_events_url": "https://api.github.com/users/octocat/received_events",
      "type": "User",
      "site_admin": false
    },
    "assets": [
      {
        "url": "https://api.github.com/repos/octocat/Hello-World/releases/assets/1",
        "browser_download_url": "https://github.com/octocat/Hello-World/releases/download/v1.0.0/example.zip",
        "id": 1,
        "node_id": "MDEyOlJlbGVhc2VBc3NldDE=",
        "name": "example.zip",
        "label": "short description",
        "state": "uploaded",
        "content_type": "application/zip",
        "size": 1024,
        "download_count": 42,
        "created_at": "2013-02-27T19:35:32Z",
        "updated_at": "2013-02-27T19:35:32Z",
        "uploader": {
          "login": "octocat",
          "id": 1,
          "node_id": "MDQ6VXNlcjE=",
          "avatar_url": "https://github.com/images/error/octocat_happy.gif",
          "gravatar_id": "",
          "url": "https://api.github.com/users/octocat",
          "html_url": "https://github.com/octocat",
          "followers_url": "https://api.github.com/users/octocat/followers",
          "following_url": "https://api.github.com/users/octocat/following{/other_user}",
          "gists_url": "https://api.github.com/users/octocat/gists{/gist_id}",
          "starred_url": "https://api.github.com/users/octocat/starred{/owner}{/repo}",
          "subscriptions_url": "https://api.github.com/users/octocat/subscriptions",
          "organizations_url": "https://api.github.com/users/octocat/orgs",
          "repos_url": "https://api.github.com/users/octocat/repos",
          "events_url": "https://api.github.com/users/octocat/events{/privacy}",
          "received_events_url": "https://api.github.com/users/octocat/received_events",
          "type": "User",
          "site_admin": false
        }
      }
    ]
  }
]
 */


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

#[derive(serde::Deserialize, Debug)]
struct GithubRelease {
    url: String,
    html_url: String,
    assets_url: String,
    upload_url: String,
    tarball_url: String,
    zipball_url: String,
    id: i64,
    node_id: String,
    tag_name: String,
    target_commitish: String,
    name: String,
    body: String,
    draft: bool,
    prerelease: bool,
    created_at: String,
    published_at: String,
    author: GithubUser,
    assets: Vec<GithubReleaseAsset>
}

#[derive(serde::Deserialize, Debug)]
struct GithubUser {
    login: String,
    id: i64,
    node_id: String,
    avatar_url: String,
    gravatar_id: String,
    url: String,
    html_url: String,
    followers_url: String,
    following_url: String,
    gists_url: String,
    starred_url: String,
    subscriptions_url: String,
    organizations_url: String,
    repos_url: String,
    events_url: String,
    received_events_url: String,
    r#type: String,
    site_admin: bool
}

#[derive(serde::Deserialize, Debug)]
struct GithubReleaseAsset {
    url: String,
    browser_download_url: String,
    id: i64,
    node_id: String,
    name: String,
    label: String,
    state: String,
    content_type: String,
    size: i64,
    download_count: i64,
    created_at: String,
    updated_at: String,
    uploader: GithubUser
}

fn main() {
    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" => "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" => "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");

    // get version info and codegen

    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
}