path: root/src/main.rs

#[macro_use] extern crate serde_derive;

use std::{
    env, iter,
    fs::File,
    io::{ self, Read, BufRead, BufReader, Write },
    path::PathBuf
};

use config::{Config, ConfigBuilder};
use structopt::StructOpt;

use mime_guess::{ Mime, MimeGuess };

use bytes::Bytes;
use futures::stream::{ self, StreamExt };
use rusoto_core::{
    request,
    Region, ByteStream,
    credential::StaticProvider
};
use rusoto_s3::{
    S3, S3Client,
    PutObjectRequest,
    HeadObjectRequest
};

mod base62;

const ONE_TRUE_THEME: &str = "Solarized (dark)";

#[derive(StructOpt, Debug, Clone, Serialize, Deserialize)]
struct Opt {
    #[structopt(short = "C", long = "config-file", parse(from_os_str))]
    config_file: Option<PathBuf>,

    /// Read from given file, instead of stdin
    #[structopt(short, long, parse(from_os_str))]
    file: Option<PathBuf>,

    /// Highlight as code
    #[structopt(short, long)]
    code: bool,

    /// Override language used to highlight input
    #[structopt(short, long)]
    language: Option<String>,

    /// Set name for paste
    #[structopt(short, long)]
    name: Option<String>,

    /// Size of random identifier: will be chosen in 0..62^L
    #[structopt(short = "L", long, default_value="3")]
    id_length: u32,

    /// Generate longer identifier, which is harder to brute-force
    #[structopt(short, long)]
    private: bool,

    /// Override mimetype for paste
    #[structopt(short, long)]
    mime: Option<String>,
}

fn main() {
    let cfg = {
        let opt = Opt::from_args();
        let mut cfg = ConfigBuilder::<config::builder::DefaultState>::default();

        let xdg = xdg::BaseDirectories::with_prefix(env!("CARGO_PKG_NAME")).unwrap();
        let xdg_file = xdg.find_config_file("config.json");
        let cfg_file = opt.config_file.clone()
            .or_else(|| env::var("PB_CONFIG_FILE").map(PathBuf::from).ok())
            .or(xdg_file);
        if let Some(cfg_file) = cfg_file {
            let cfg_file = cfg_file.to_str().expect("Config file has invalid path");
            cfg = cfg.add_source(config::File::new(cfg_file, config::FileFormat::Json));
        }
				let mut cfg = cfg
					.add_source(config::Environment::with_prefix("PB"))
					.add_source(Config::try_from(&opt).unwrap())
					.build()
					.unwrap();

        if cfg.get_bool("private").unwrap_or(false) {
            cfg.set("id_length", 10).unwrap();
        }
        cfg
    };

    smol::block_on(async_compat::Compat::new(async_main(cfg)));
}

async fn async_main(cfg: Config) {
    let region = Region::Custom {
        endpoint: {
            // trailing / causes problems with rusoto
            let mut s = cfg.get_string("endpoint").expect("No endpoint set");
            if s.ends_with('/') { s.pop(); }
            s
        },
        name: cfg.get_string("region").expect("No region set")
    };

    let path = cfg.get_string("file");
    let lang = cfg.get_string("language").ok();
    let code = cfg.get_bool("code").unwrap()
        || lang.is_some();

    let input: Box<dyn BufRead + Send + Sync> = if let Ok(ref path) = path {
        Box::new(BufReader::new(File::open(path).unwrap()))
    } else {
        let stdin = io::stdin();
        Box::new(BufReader::new(stdin))
    };

    let data: Box<dyn Iterator<Item=Bytes> + Send + Sync> =
        if code { Box::new(highlight(&cfg, input)) }
        else { Box::new(chunk_iter(input)) };

    let body = ByteStream::new(stream::iter(data).map(Ok));

    let dispatcher = request::HttpClient::new().expect("Unable to create rusoto http client");
    let credentials = StaticProvider::new_minimal(
        cfg.get_string("access_key_id").expect("Access key not set"),
        cfg.get_string("secret_access_key").expect("Secret access key not set"));
    let client = S3Client::new_with(dispatcher, credentials, region);

    let bucket = cfg.get_string("bucket").expect("Bucket not set");
    let name = cfg.get_string("name").ok()
        .unwrap_or_else(|| {
            let len = cfg.get_int("id_length").unwrap() as u32;
            (0..).map(|_| random_id(len))
                 .find(|key| !smol::block_on(check_exists(&client, bucket.clone(), key.clone())))
                 .unwrap()
        });

    let mime = if code { String::from("text/html; charset=utf-8") } else {
        cfg.get_string("mime").ok()
           .or_else(|| path.ok()
                        .and_then(|path| MimeGuess::from_path(path).first())
                        .as_ref()
                        .map(Mime::to_string))
           .unwrap_or_else(|| String::from("text/plain; charset=utf-8"))
    };

    let _put = client.put_object(PutObjectRequest {
        bucket,
        key: name.clone(),
        content_type: Some(mime),
        body: Some(body),
        ..Default::default()
    }).await.expect("Put failed");

    println!("{}", name);
}

static TEMPLATE: &str = include_str!("./code.html");
#[derive(Serialize)]
struct Template { title: String, code: String }
fn highlight(cfg: &Config, mut input: Box<dyn BufRead + Send>) -> impl Iterator<Item=Bytes> {
    use syntect::{
        parsing::{ SyntaxSet, SyntaxDefinition },
        highlighting::ThemeSet,
        easy::HighlightLines,
        html::{ start_highlighted_html_snippet, append_highlighted_html_for_styled_line, IncludeBackground },
        util::LinesWithEndings
    };

    use tinytemplate::TinyTemplate;

    let mut input_str = String::new();
    input.read_to_string(&mut input_str).expect("Unable to read input");

    let mut syntax_set = SyntaxSet::load_defaults_newlines().into_builder();
    syntax_set.add(SyntaxDefinition::load_from_str(include_str!("../syntax/nix.sublime-syntax"), false, None)
                    .unwrap());
    let syntax_set = syntax_set.build();

    let syntax =
        if let Ok(lang) = cfg.get_string("language") {
            syntax_set.find_syntax_by_token(&lang).expect("Failed to load syntax")
        } else if let Ok(path) = cfg.get_string("file") {
            syntax_set.find_syntax_for_file(path)
                .unwrap_or_default().expect("Failed to load syntax from path")
        } else {
            syntax_set.find_syntax_by_first_line(&input_str)
                      .unwrap_or_else(|| syntax_set.find_syntax_plain_text())
        };

    let theme_set = ThemeSet::load_defaults();
    let theme = theme_set.themes.get(ONE_TRUE_THEME).expect("Failed to load theme");

    let mut highlighter = HighlightLines::new(syntax, theme);

    let mut html = Vec::new();
    let mut formatted_line = String::new();
    let (_, bg) = start_highlighted_html_snippet(theme);
    let _ = write!(&mut html, "<table><tr><td><pre class=\"linenumbers\" disabled>");
    for (i, _line) in input_str.lines().enumerate() {
        let _ = write!(&mut html, "<a href=\"#L{}\" data-txt=\"{}\"></a>", i, i);
    }
    let _ = write!(&mut html, "</pre></td><td><pre class=code><code>");

    for (i, line) in LinesWithEndings::from(&input_str).enumerate() {
        let regions = highlighter.highlight(line, &syntax_set);
        formatted_line.clear();
        append_highlighted_html_for_styled_line(&regions[..], IncludeBackground::IfDifferent(bg), &mut formatted_line).unwrap();
        if formatted_line.ends_with('\n') { formatted_line.pop(); }
        let _ = write!(&mut html, "<a id=\"L{}\"></a>{}", i, formatted_line);
    }
    let _ = write!(&mut html, "</code></pre></td></tr></table>");

    let mut template = TinyTemplate::new();
    template.add_template("code", TEMPLATE).expect("Template is faulty");
    template.add_formatter("raw", tinytemplate::format_unescaped);

    let rendered = template.render("code", &Template {
        title: cfg.get_string("name").unwrap_or_else(|_| String::new()),
        code: String::from_utf8(html).expect("Invalid UTF8")
    }).unwrap();

    iter::once(Bytes::copy_from_slice(rendered.as_bytes()))
}

fn chunk_iter(input: Box<dyn BufRead + Send + Sync>) -> impl Iterator<Item=Bytes> {
    itertools::unfold(input, |input| {
        let buf = match input.fill_buf().ok() {
            Some([]) => None,
            Some(buf) => Some(Bytes::copy_from_slice(buf)),
            None => None
        };

        if let Some(buf) = &buf { input.consume(buf.len()) }
        buf
    })
}

async fn check_exists(c: &S3Client, bucket: String, key: String) -> bool {
    c.head_object(HeadObjectRequest {
        bucket, key,
        ..Default::default()
    }).await.is_ok()
}

fn random_id(size: u32) -> String {
    use rand::distributions::{ Distribution, Uniform };
    use base62::*;
    let range = Uniform::from(0..62u128.pow(size));
    let n = range.sample(&mut rand::thread_rng());
    digits_to_string(number_to_digits(n, BASE).into_iter(), ALPHABET)
}