// This file is part of HTTP Signatures

// HTTP Signatures is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// HTTP Signatures is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with HTTP Signatures  If not, see <http://www.gnu.org/licenses/>.

use std::io::Read;
use std::convert::{TryFrom, TryInto};
use std::collections::HashMap;

use ring::{signature, digest, hmac};
use ring::error::Unspecified;
use base64::decode;
use untrusted::Input;

use super::{SignatureAlgorithm, ShaSize, REQUEST_TARGET};
use error::{DecodeError, VerificationError};

const KEY_ID: &'static str = "keyId";
const HEADERS: &'static str = "headers";
const ALGORITHM: &'static str = "algorithm";
const DATE: &'static str = "date";
const SIGNATURE: &'static str = "signature";

/// The GetKey trait is used during HTTP Signature verification to access the required decryption
/// key based on a given key_id.
///
/// The `key_id` is provided in the Authorization header of the request as `KeyId`.
///
/// ### Example
/// ```rust
/// # use std::io::Cursor;
/// # use std::collections::HashMap;
/// use http_signatures::GetKey;
///
/// struct MyKeyGetter {
///     keys: HashMap<String, Vec<u8>>,
/// }
///
/// impl MyKeyGetter {
///     pub fn new() -> Self {
///         MyKeyGetter {
///             keys: HashMap::new(),
///         }
///     }
///
///     pub fn add_key(&mut self, key_id: String, key: Vec<u8>) {
///         self.keys.insert(key_id, key);
///     }
/// }
///
/// impl GetKey for MyKeyGetter {
///     type Key = Cursor<Vec<u8>>;
///     type Error = ();
///
///     fn get_key(self, key_id: &str) -> Result<Self::Key, Self::Error> {
///         self.keys.get(key_id).map(|key| Cursor::new(key.clone())).ok_or(())
///     }
/// }
///
/// # fn run() -> Result<(), ()> {
/// let mut key_getter = MyKeyGetter::new();
/// key_getter.add_key("key-1".into(), vec![1, 2, 3, 4, 5]);
///
/// key_getter.get_key("key-1")?;
/// # Ok(())
/// # }
/// ```
pub trait GetKey {
    type Key: Read;
    type Error;

    fn get_key(self, key_id: &str) -> Result<Self::Key, Self::Error>;
}

/// The `VerifyAuthorizationHeader` trait is meant to be implemented for the request types from
/// http libraries (such as Hyper and Rocket). This trait makes verifying requests much easier,
/// since the `verify_authorization_header()` method can be called directly on a Request type.
///
/// For examples, see the
/// [hyper server](https://github.com/asonix/http-signatures/blob/master/examples/hyper_server.rs)
/// and [rocket](https://github.com/asonix/http-signatures/blob/master/examples/rocket.rs) files.
pub trait VerifyAuthorizationHeader {
    fn verify_authorization_header<G: GetKey>(
        &self,
        key_getter: G,
    ) -> Result<(), VerificationError>;
}

/// The `AuthorizationHeader` struct is the direct reasult of reading in the Authorization header
/// from a given request.
///
/// It contains the keys to the request's headers in the correct order for recreating the signing
/// string, the algorithm used to create the signature, and the signature itself.
///
/// It also contains the key_id, which will be handled by a type implementing `GetKey`.
pub struct AuthorizationHeader<'a> {
    key_id: &'a str,
    header_keys: Vec<&'a str>,
    algorithm: SignatureAlgorithm,
    signature: Vec<u8>,
}

impl<'a> AuthorizationHeader<'a> {
    /// Try to create an `AuthorizationHeader` from a given String.
    pub fn new(s: &'a str) -> Result<Self, DecodeError> {
        s.try_into()
    }

    /// Try to verify the current `AuthorizationHeader`.
    pub fn verify<G>(
        self,
        headers: &[(&str, &str)],
        method: &str,
        path: &str,
        query: Option<&str>,
        key_getter: G,
    ) -> Result<(), VerificationError>
    where
        G: GetKey,
    {
        let vah = CheckAuthorizationHeader {
            auth_header: self,
            headers: headers,
            method: method,
            path: path,
            query: query,
        };

        vah.verify(key_getter)
    }
}

impl<'a> TryFrom<&'a str> for AuthorizationHeader<'a> {
    type Error = DecodeError;

    fn try_from(s: &'a str) -> Result<Self, Self::Error> {
        let s = s.trim_left_matches("Signature: ");
        let key_value = s.split(',')
            .filter_map(|item| {
                let eq_index = item.find("=")?;
                let tup = item.split_at(eq_index);
                let val = tup.1.get(1..)?;
                Some((tup.0, val))
            })
            .collect::<HashMap<&str, &str>>();

        let key_id = key_value
            .get(KEY_ID)
            .ok_or(DecodeError::MissingKey(KEY_ID))?
            .trim_left_matches("\"")
            .trim_right_matches("\"");

        let header_keys = key_value
            .get(HEADERS)
            .unwrap_or(&DATE)
            .trim_left_matches("\"")
            .trim_right_matches("\"")
            .split(' ')
            .collect();

        let algorithm = (*key_value
                             .get(ALGORITHM)
                             .ok_or(DecodeError::MissingKey(ALGORITHM))?
                             .trim_left_matches("\"")
                             .trim_right_matches("\""))
            .try_into()?;

        let sig_string: String = key_value
            .get(SIGNATURE)
            .ok_or(DecodeError::MissingKey(SIGNATURE))?
            .trim_left_matches("\"")
            .trim_right_matches("\"")
            .into();

        let signature = decode(&sig_string).map_err(|_| DecodeError::NotBase64)?;

        Ok(AuthorizationHeader {
            key_id,
            header_keys,
            algorithm,
            signature,
        })
    }
}

struct CheckAuthorizationHeader<'a> {
    auth_header: AuthorizationHeader<'a>,
    headers: &'a [(&'a str, &'a str)],
    method: &'a str,
    path: &'a str,
    query: Option<&'a str>,
}

impl<'a> CheckAuthorizationHeader<'a> {
    pub fn verify<G>(&self, key_getter: G) -> Result<(), VerificationError>
    where
        G: GetKey,
    {
        let key: G::Key = key_getter.get_key(self.auth_header.key_id).map_err(|_| {
            VerificationError::GetKey
        })?;

        let headers: HashMap<String, Vec<&str>> =
            self.headers.iter().fold(HashMap::new(), |mut acc,
             &(ref key, ref value)| {
                acc.entry(key.to_lowercase()).or_insert(Vec::new()).push(
                    value,
                );

                acc
            });

        let mut headers: HashMap<&str, String> = headers
            .iter()
            .map(|(key, value)| (key.as_ref(), value.join(", ")))
            .collect();

        headers.insert(
            REQUEST_TARGET.into(),
            if let Some(ref query) = self.query {
                format!(
                    "{} {}?{}",
                    self.method.to_lowercase(),
                    self.path,
                    query,
                )
            } else {
                format!(
                    "{} {}",
                    self.method.to_lowercase(),
                    self.path,
                )
            },
        );

        let signing_vec = self.auth_header.header_keys.iter().fold(
            (Vec::new(), Vec::new()),
            |mut acc, header_key| {
                if let Some(ref header) = headers.get(header_key) {
                    acc.0.push(format!("{}: {}", header_key, header));
                } else {
                    acc.1.push(header_key.clone());
                }

                acc
            },
        );

        if signing_vec.1.len() > 0 {
            return Err(VerificationError::MissingHeaders(signing_vec.1.join(", ")));
        }

        let signing_string = signing_vec.0.join("\n");

        match self.auth_header.algorithm {
            SignatureAlgorithm::RSA(ref sha_size) => {
                Self::verify_rsa(
                    key,
                    sha_size,
                    signing_string,
                    self.auth_header.signature.as_ref(),
                )
            }
            SignatureAlgorithm::HMAC(ref sha_size) => {
                Self::verify_hmac(
                    key,
                    sha_size,
                    signing_string,
                    self.auth_header.signature.as_ref(),
                )
            }
        }
    }

    fn verify_rsa<T>(
        mut key: T,
        sha_size: &ShaSize,
        signing_string: String,
        sig: &[u8],
    ) -> Result<(), VerificationError>
    where
        T: Read,
    {
        // Verify the signature.
        let mut public_key_der = Vec::new();
        key.read_to_end(&mut public_key_der).map_err(|_| {
            VerificationError::ReadKey
        })?;
        let public_key_der = Input::from(&public_key_der);
        let message = Input::from(signing_string.as_ref());
        let signature = Input::from(sig);

        match *sha_size {
            ShaSize::TwoFiftySix => {
                signature::verify(
                    &signature::RSA_PKCS1_2048_8192_SHA256,
                    public_key_der,
                    message,
                    signature,
                ).map_err(|Unspecified| VerificationError::BadSignature)?;
            }
            ShaSize::ThreeEightyFour => {
                signature::verify(
                    &signature::RSA_PKCS1_2048_8192_SHA384,
                    public_key_der,
                    message,
                    signature,
                ).map_err(|Unspecified| VerificationError::BadSignature)?;
            }
            ShaSize::FiveTwelve => {
                signature::verify(
                    &signature::RSA_PKCS1_2048_8192_SHA512,
                    public_key_der,
                    message,
                    signature,
                ).map_err(|Unspecified| VerificationError::BadSignature)?;
            }
        }

        Ok(())
    }

    fn verify_hmac<T>(
        mut key: T,
        sha_size: &ShaSize,
        signing_string: String,
        sig: &[u8],
    ) -> Result<(), VerificationError>
    where
        T: Read,
    {
        let mut hmac_key = Vec::new();
        key.read_to_end(&mut hmac_key).map_err(
            |_| VerificationError::ReadKey,
        )?;
        let hmac_key = hmac::SigningKey::new(
            match *sha_size {
                ShaSize::TwoFiftySix => &digest::SHA256,
                ShaSize::ThreeEightyFour => &digest::SHA384,
                ShaSize::FiveTwelve => &digest::SHA512,
            },
            &hmac_key,
        );

        hmac::verify_with_own_key(&hmac_key, signing_string.as_ref(), sig)
            .map_err(|_| VerificationError::BadSignature)?;

        Ok(())
    }
}