#![feature(proc_macro_hygiene, decl_macro)]

#[macro_use] extern crate rocket;

use std::fmt;
use std::io;

use rocket_contrib::templates::Template;
use rocket_contrib::helmet::*;

use rocket::Request;
use rocket::request;
use rocket::request::{FromRequest,FromFormValue,LenientForm};
use rocket::response::NamedFile;
use rocket::outcome::Outcome;
use rocket::http::{Status,RawStr};
use rocket::State;

use serde::Serialize;
use serde::Serializer;

use thiserror::Error;
use anyhow::Context;

//use std::sync::Arc;
use std::rc::Rc;
use std::fs;
//use std::ops::Deref;

type E = anyhow::Error;

mod ticket {
  use base64::{URL_SAFE,URL_SAFE_NO_PAD};
  use thiserror::Error;
  use tokio::time::{Instant,Duration};
  use std::fs::File;
  use std::io::Read;
  use std::convert::TryInto;
  use std::array::TryFromSliceError;
  use hmac::{Mac,Hmac};
  use generic_array::GenericArray;

  #[derive(Error,Debug)]
  pub enum DecapError {
    #[error("bad ticket format")]
    BadFormat,
    #[error("stale ticket?")]
    Stale,
    #[error("ticket out of sequence!")]
    BadSequence,
  }

  impl From<base64::DecodeError> for DecapError {
    fn from(_ : base64::DecodeError) -> Self { BadFormat }
  }
  impl From<TryFromSliceError> for DecapError {
    fn from(_ : TryFromSliceError) -> Self { BadFormat }
  }

  use DecapError::*;
  type HmacAlgo = Hmac<sha2::Sha512Trunc224>;

  pub struct Key {
    key : GenericArray<u8, <HmacAlgo as Mac>::KeySize>,
    offset : u128,
    epoch : Instant,
  }
  impl Key {
    pub fn new() -> Result<Key,anyhow::Error> {
      let mut key = <GenericArray<u8,_> as Default>::default();
      let mut f = File::open("/dev/urandom")?;
      f.read_exact(&mut key)?;
      let mut offset : [u8;16] = Default::default();
      f.read_exact(&mut offset)?;
      let offset = u128::from_le_bytes(offset);
      Ok(Key{ key, offset, epoch : Instant::now() })
    }
    fn encapsulate(&self, msg : &[u8]) -> String {
      let mut maccer = HmacAlgo::new(&self.key);
      maccer.input(&msg);
      let tag = maccer.result().code();
      let msg = base64::encode_config(msg,  URL_SAFE);
      let tag = base64::encode_config(&tag, URL_SAFE_NO_PAD);
      format!("{}.{}",msg,tag)
    }
    fn decapsulate(&self, both : &str) -> Result<Vec<u8>,DecapError> {
      let mut split = both.splitn(2,".");
      let msg = split.next().ok_or(BadFormat)?;
      let gottag = split.next().ok_or(BadFormat)?;
      let msg = base64::decode_config(&msg, URL_SAFE_NO_PAD)?;
      let mut gottagbuf = <GenericArray<u8,_> as Default>::default();
      base64::decode_config_slice(&gottag, URL_SAFE_NO_PAD, &mut gottagbuf)?;
      let gottag = gottagbuf;

      let mut maccer = HmacAlgo::new(&self.key);
      maccer.input(&msg);
      let exptag = maccer.result().code();
      if exptag != gottag { return Err(Stale) }
      Ok(msg)
    }

    pub fn make_ticket(&self) -> String {
      let now = self.epoch.elapsed().as_millis();
      let msg = u128::to_le_bytes(now + self.offset);
      let msg = self.encapsulate(&msg);
      msg
    }

    pub fn decode_ticket(&self, both : &str) -> Result<Duration,DecapError> {
      let msg = self.epoch.elapsed().as_millis();
      let now = msg - self.offset;
      let then = self.decapsulate(&both)?;
      let then : &[u8] = &then;
      let then = u128::from_le_bytes(then.try_into()?);
      if then > now { Err(BadSequence)? }
      Ok(Duration::from_millis((now - then).try_into()
                               .map_err(|_| Stale)?))
    }
  }
}

struct WebUiState {
  #[allow(dead_code)]
  cfg : picollar::config::Config,
  magic : String,
  template_dir : String,
  ticket_key : ticket::Key,
}

#[derive(Error,Debug,Eq,PartialEq)]
enum MagicError {
  #[error("missing magic")]
  Missing,
  #[error("wrong magic")]
  Wrong,
  #[error("badly formatted parameter for magic")]
  ParseError,
}

struct MagicOk (Rc<String>);
impl Serialize for MagicOk {
  fn serialize<S>(&self, ser : S) -> Result<S::Ok, S::Error> where S : Serializer {
    self.0.serialize(ser)
  }
}

impl fmt::Debug for MagicOk {
  fn fmt(&self, fmt : &mut fmt::Formatter) -> Result<(),fmt::Error> {
    write!(fmt,"MagicOk(/*elided*/)")
  }
}

impl MagicOk {
  fn from_result<E>(st : &WebUiState, rs : Result<&UncheckedMagic, E>)
                    -> Result<MagicOk,MagicError> {
    match rs {
      Ok(s) => MagicOk::new(st, s),
      Err(_e) => Err(MagicError::ParseError),
    }
  }
  fn new(st : &WebUiState, um : &UncheckedMagic)
         -> Result<MagicOk,MagicError> {
    if *um.danger == st.magic { Ok(MagicOk(um.danger.clone())) }
    else { Err(MagicError::Wrong) }
  }
}

impl<'a, 'r> FromRequest<'a, 'r> for MagicOk {
  type Error = MagicError;
  fn from_request(req : &'a Request<'r>)
                  -> request::Outcome<MagicOk, MagicError> {
    let st = req.guard::<State<WebUiState>>().expect("get state");
    use Outcome::*;
    use MagicError::*;
    match req.get_query_value("m") {
      Some(rs) => match MagicOk::from_result(&st, rs.as_ref()) {
        Ok(m) => Success(m),
        Err(e) => Failure((Status::Forbidden, e)),
      },
      None => return Failure((Status::Forbidden, Missing)),
    }
  }
}

struct UncheckedMagic { danger : Rc<String> }

impl<'v> FromFormValue<'v> for UncheckedMagic {
  type Error = MagicError;
  fn from_form_value(v: &'v RawStr) -> Result<UncheckedMagic, MagicError> {
    let rs = v.parse().map_err(|_e| MagicError::ParseError)?;
    Ok(UncheckedMagic { danger : Rc::new(rs) })
  }
}

#[derive(Serialize,Debug)]
struct MinimalRenderContext { m : MagicOk }

#[get("/")]
fn welcome(m : MagicOk) -> Template {
  let context = MinimalRenderContext{ m : m };
  Template::render("welcome",&context)
}

#[derive(FromForm)]
struct MainForm {
  m : UncheckedMagic,
}

#[post("/main", data="<data>")]
fn main_screen(data : LenientForm<MainForm>,
               st : State<WebUiState>) -> Result<Template,MagicError> {
  let c = MinimalRenderContext{ m : MagicOk::new(&st, &data.m)? };
  Ok(Template::render("main",&c))
}  

#[post("/api/ticket")]
fn get_ticket(st : State<WebUiState>) -> Result<String,MagicError> {
  Ok( st.ticket_key.make_ticket() )
}  

#[get("/script.js")]
fn script(st : State<WebUiState>) -> io::Result<NamedFile> {
  NamedFile::open(format!("{}/script.js", st.template_dir))
}  

fn main() -> Result<(),E> {
  let cfg = picollar::config::read()?;
  let template_dir : String = cfg.need(&["web_ui","template_dir"])?;

  let helmet = SpaceHelmet::default()
    .enable(NoSniff::Enable)
    .enable(Frame::Deny)
    .enable(Referrer::NoReferrer);

  let mut cbuilder = rocket::config::Config::build(
    rocket::config::Environment::Staging
  )
    .address("127.0.0.1")
    .port(16714)
    .workers(50)
    .keep_alive(20);
  cbuilder
    .extras.insert("template_dir".to_owned(), template_dir.clone().into());
  let rconfig = cbuilder
    .finalize()?;

  let magic_file : String = cfg.need(&["web_ui","magic_file"])?;
  let magic = fs::read_to_string(magic_file).context("magic file")?;
  let magic = magic.trim().to_owned();

  let ticket_key = ticket::Key::new()?;

  rocket::custom(rconfig)
    .attach(helmet)
    .attach(Template::fairing())
    .manage(WebUiState { cfg, template_dir, magic, ticket_key })
    .mount("/", routes![
      welcome,
      main_screen,
      script,
      get_ticket,
    ]).launch();

  Ok(())
}