sofapub 0.1.9

use std::collections::HashMap;
use rocket::{
    http::Status,
    request::{FromRequest, Outcome, Request},
};
use base64::{engine::general_purpose, engine::Engine as _};
use rsa::signature::Verifier;
use rsa::RsaPublicKey;
use rsa::pkcs8::DecodePublicKey;

use sha2::Sha256;
use serde_json::Value;

use crate::{Configuration, server::ManagedConfiguration};
use crate::signing::maybe_signed_get;

#[derive(Clone, Debug)]
pub struct VerifyParams {
    pub signature: String,
    pub request_target: String,
    pub host: String,
    pub date: String,
    pub digest: Option<String>,
    pub content_type: String,
    pub user_agent: Option<String>,
}

fn build_verify_string(params: VerifyParams) -> (String, String, String) {
    let mut signature_map = HashMap::<String, String>::new();
    let parts_re = regex::Regex::new(r#"(\w+)="(.+?)""#).expect("Invalid regex pattern");

    for cap in parts_re.captures_iter(&params.signature) {
        signature_map.insert(cap[1].to_string(), cap[2].to_string());
    }

    let key_id = signature_map.get("keyId").expect("keyId not found in signature_map");
    let key_id_parts: Vec<_> = key_id.split('#').collect();
    let ap_id = key_id_parts.first().expect("Failed to parse ap_id").to_string();

    let headers = signature_map
        .get("headers")
        .expect("headers not found in signature_map");

    let verify_string = headers
        .split_whitespace()
        .map(|part| match part {
            "(request-target)" => format!("(request-target): {}", params.request_target),
            "host" => format!("host: {}", params.host),
            "date" => format!("date: {}", params.date),
            "digest" => format!("digest: {}", params.digest.clone().unwrap_or_default()),
            "content-type" => format!("content-type: {}", params.content_type),
            "user-agent" => format!("user-agent: {}", params.user_agent.clone().unwrap_or_default()),
            _ => String::new(),
        })
        .collect::<Vec<String>>()
        .join("\n");

    (
        verify_string,
        signature_map.get("signature").expect("signature not found in signature_map").clone(),
        ap_id,
    )
}

pub async fn verify(configuration: Configuration, params: VerifyParams) -> bool {
    let (verify_string, signature_str, ap_id) = build_verify_string(params);

    async fn fetch_actor(configuration: Configuration, ap_id: &str) -> Option<Value> {
        maybe_signed_get(configuration, ap_id.to_string(), false).await.ok()?.json::<Value>().await.ok()
    }

    fn verify_signature(public_key: RsaPublicKey, signature_str: &str, verify_string: &str) -> bool {
        let verifying_key = rsa::pkcs1v15::VerifyingKey::<Sha256>::new(public_key);
        let signature_bytes = general_purpose::STANDARD.decode(signature_str.as_bytes()).unwrap();
        let signature = rsa::pkcs1v15::Signature::try_from(signature_bytes.as_slice()).unwrap();

        verifying_key.verify(verify_string.as_bytes(), &signature).is_ok()
    }

    fetch_actor(configuration, &ap_id).await.and_then(|actor| {
        actor.get("publicKey")
            .and_then(|pk| pk.get("publicKeyPem"))
            .and_then(|pkp| RsaPublicKey::from_public_key_pem(pkp.as_str().unwrap()).ok())
            .map(|public_key| verify_signature(public_key, &signature_str, &verify_string))
    }).unwrap_or_else(|| {
        log::debug!("failed to retrieve or deserialize actor");
        false
    })
}

pub struct Signed(pub bool);

#[derive(Debug)]
pub enum SignatureError {
    NonExistent,
    MultipleSignatures,
    InvalidRequestPath,
    InvalidRequestUsername,
    LocalUserNotFound,
    SignatureInvalid,
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for Signed {
    type Error = SignatureError;

    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        let configuration = request.guard::<ManagedConfiguration>().await.unwrap();
        
        let request_target = format!("{} {}", request.method().to_string().to_lowercase(), request.uri().to_string().trim_end_matches('&'));

        fn fetch_header_value(request: &Request, header_name: &str, backup_name: Option<&str>) -> Option<String> {
            let mut values: Vec<_> = request.headers().get(header_name).collect();

            if values.len() != 1 {
                if let Some(backup) = backup_name {
                    values = request.headers().get(backup).collect();
                }
            }

            values.first().map(|s| s.to_string())
        }

        // browser fetch is a jerk and forbids the "date" header; browsers
        // aggressively strip it, so I use Enigmatick-Date as a backup
        let date = fetch_header_value(request, "date", Some("enigmatick-date"));
        let digest = fetch_header_value(request, "digest", None);
        let user_agent = fetch_header_value(request, "user-agent", None);

        match (request.content_type(), fetch_header_value(request, "signature", None)) {
            (Some(content_type), Some(signature)) => {
                let verified = verify(
                    configuration.configuration,
                    VerifyParams {
                        signature,
                        request_target,
                        host: request.host().unwrap().to_string(),
                        date: date.unwrap_or_default(),
                        digest,
                        content_type: content_type.to_string(),
                        user_agent,
                    },
                ).await;

                Outcome::Success(Signed(verified))
            }
            (Some(_), None) => Outcome::Success(Signed(false)),
            (_, _) => Outcome::Failure((Status::BadRequest, SignatureError::NonExistent)),
        }
    }
}