#![feature(proc_macro_hygiene, decl_macro)]

#[macro_use]
extern crate rocket;

use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};

use anyhow::{anyhow, bail, Result};
use ldap3::{
    controls::{MakeCritical, RelaxRules},
    LdapConn, Mod, Scope, SearchEntry,
};
use lettre::{SmtpClient, Transport};
use lettre_email::EmailBuilder;
use log::{debug, error, info};
use rand::{self, Rng};
use rocket::{
    fairing::AdHoc,
    request::{FlashMessage, Form},
    response::{Flash, Redirect},
    State,
};
use rocket_contrib::templates::Template;

#[cfg(test)]
mod tests;

struct LdapConfig {
    server: String,
    base: String,
    filter: String,
    bind: String,
    password: String,
}

struct Ldap0rConfig {
    domain: String,
    ldap: LdapConfig,
}

const BASE62: &[u8] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";

type Keys = Arc<Mutex<HashMap<String, String>>>;

fn reset_prepare(config: &Ldap0rConfig, keys: &Keys, email_address: &str) -> Result<()> {
    // ldap lookup
    let mut ldap = LdapConn::new(&config.ldap.server)?;
    let result = ldap.search(
        &config.ldap.base,
        Scope::Subtree,
        &format!("(&{}(mail={}))", config.ldap.filter, email_address),
        vec![""],
    )?;
    ldap.unbind()?;
    let (rs, _res) = result.success()?;

    // check for less or more than 1 result
    if rs.len() != 1 {
        bail!("Invalid password reset request for '{}'", email_address);
    }

    // generate key
    let mut rng = rand::thread_rng();
    let key: String = (0..64)
        .map(|_| BASE62[rng.gen::<usize>() % 62] as char)
        .collect();

    // store key with id
    let keys = Arc::clone(&keys);
    let mut keys = keys
        .lock()
        .map_err(|e| anyhow!("Could not aquire lock for keys: {}", e))?;
    keys.insert(key.to_string(), email_address.to_string());
    debug!("Generated new key for '{}'", email_address);

    // generate email
    let email = EmailBuilder::new()
        .to(email_address)
        .from("ldap0r@example.com")
        .subject("LDAP password reset")
        .text(format!(
            "Use following url to set a new password: {}/reset/{}",
            config.domain, key
        ))
        .build()?;

    // send email
    let mut mailer = SmtpClient::new_unencrypted_localhost()?.transport();
    mailer.send(email.into())?;
    info!("Password reset email was sent to '{}'", email_address);

    Ok(())
}

fn set_password(
    config: &Ldap0rConfig,
    keys: &Keys,
    key: &str,
    passwords: &PasswordsForm,
) -> Result<Flash<Redirect>> {
    if passwords.password != passwords.password_control {
        return Ok(Flash::error(
            Redirect::to(uri!(reset_key: key)),
            "Password does not match the password verification field",
        ));
    }

    if passwords.password.len() < 8 {
        return Ok(Flash::error(
            Redirect::to(uri!(reset_key: key)),
            "Password length has to be at least 8",
        ));
    }

    // key lookup
    let keys = Arc::clone(&keys);
    let mut keys = match keys.lock() {
        Ok(keys) => keys,
        Err(e) => {
            error!("Could not aquire lock for keys: {}", e);
            return Ok(Flash::error(
                Redirect::to(uri!(reset_key: key)),
                "Setting new password failed",
            ));
        }
    };
    let email = keys
        .get(key)
        .ok_or_else(|| anyhow!("Could not extract email"))?
        .to_string();

    // ldap lookup
    let mut ldap = LdapConn::new(&config.ldap.server)?;
    let result = ldap.search(
        &config.ldap.base,
        Scope::Subtree,
        &format!("(&{}(mail={}))", &config.ldap.filter, &email),
        vec!["cn"],
    )?;
    let (mut rs, _res) = result.success()?;

    // ldap set new password
    let user = SearchEntry::construct(
        rs.pop()
            .ok_or_else(|| anyhow!("Could extract not receive LDAP result"))?,
    )
    .attrs
    .get("cn")
    .ok_or_else(|| anyhow!("Could not extract 'cn' from LDAP entry"))?[0]
        .to_string();
    let mut password = HashSet::new();
    password.insert(passwords.password.as_str());
    ldap.simple_bind(&config.ldap.bind, &config.ldap.password)?
        .success()?;
    ldap.with_controls(RelaxRules.critical())
        .modify(
            &format!("cn={},{}", &user, &config.ldap.base),
            vec![Mod::Replace("userPassword", password)],
        )?
        .success()?;

    ldap.unbind()?;
    keys.remove(key);

    info!(
        "New password set for user '{}' with email address '{}'",
        &user, &email
    );

    Ok(Flash::success(
        Redirect::to(uri!(reset)),
        "New password was saved",
    ))
}

#[derive(Debug, FromForm)]
struct EmailForm {
    email: String,
}

#[derive(Debug, FromForm)]
struct PasswordsForm {
    password: String,
    password_control: String,
}

#[get("/")]
fn index() -> Redirect {
    Redirect::to(uri!(reset))
}

#[get("/reset")]
fn reset(flash: Option<FlashMessage>) -> Template {
    let mut context = HashMap::new();

    if let Some(ref msg) = flash {
        context.insert("flash", msg.msg());
        if msg.name() == "error" {
            context.insert("flash_type", "Error");
        }
    }

    Template::render("reset", &context)
}

#[post("/reset", data = "<email>")]
fn reset_email(
    config: State<Ldap0rConfig>,
    remote_address: SocketAddr,
    keys: State<Keys>,
    email: Form<EmailForm>,
) -> Flash<Redirect> {
    if let Err(e) = reset_prepare(&config, &keys, &email.email) {
        error!("{} from {}", e, remote_address);
    }

    Flash::success(
        Redirect::to(uri!(reset)),
        format!("Email was sent to {}", email.email),
    )
}

#[get("/reset/<key>")]
fn reset_key(keys: State<Keys>, key: String, flash: Option<FlashMessage>) -> Option<Template> {
    let keys = Arc::clone(&keys);
    if let Ok(keys) = keys.lock() {
        let mut context = HashMap::new();

        if let Some(ref msg) = flash {
            context.insert("flash", msg.msg());
            if msg.name() == "error" {
                context.insert("flash_type", "Error");
            }
        }

        if keys.contains_key(&key) {
            context.insert("key", &key);
            return Some(Template::render("reset_key", &context));
        }
    } else {
        error!("Could not aquire lock for keys");
    }

    None
}

#[post("/reset/<key>", data = "<passwords>")]
fn reset_password(
    config: State<Ldap0rConfig>,
    remote_address: SocketAddr,
    keys: State<Keys>,
    key: String,
    passwords: Form<PasswordsForm>,
) -> Flash<Redirect> {
    set_password(&config, &keys, &key, &passwords).unwrap_or_else(|e| {
        error!("{} from {}", e, remote_address);
        Flash::error(
            Redirect::to(uri!(reset_key: key)),
            "Setting new password failed",
        )
    })
}

fn main() {
    rocket::ignite()
        .mount(
            "/",
            routes![index, reset, reset_email, reset_key, reset_password],
        )
        .attach(AdHoc::on_attach("LDAP config", |rocket| {
            let ldap0r_config = Ldap0rConfig {
                domain: rocket
                    .config()
                    .get_str("domain")
                    .unwrap_or("https://localhost")
                    .to_string(),
                ldap: LdapConfig {
                    server: rocket
                        .config()
                        .get_str("ldap_server")
                        .unwrap_or("ldap://localhost:2389")
                        .to_string(),
                    base: rocket
                        .config()
                        .get_str("ldap_base")
                        .unwrap_or("ou=Places,dc=example,dc=org")
                        .to_string(),
                    filter: rocket
                        .config()
                        .get_str("ldap_filter")
                        .unwrap_or("(objectClass=person)")
                        .to_string(),
                    bind: rocket
                        .config()
                        .get_str("ldap_bind")
                        .unwrap_or("cn=admin,dc=example,dc=org")
                        .to_string(),
                    password: rocket
                        .config()
                        .get_str("ldap_password")
                        .unwrap_or("secure")
                        .to_string(),
                },
            };

            Ok(rocket.manage(ldap0r_config))
        }))
        .attach(Template::fairing())
        .manage(Keys::new(Mutex::new(HashMap::new())))
        .launch();
}