markbase/markbase-core/src/s3_policy.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BucketPolicy {
    #[serde(rename = "Version")]
    pub version: String,
    #[serde(rename = "Statement")]
    pub statement: Vec<PolicyStatement>,
}

impl Default for BucketPolicy {
    fn default() -> Self {
        Self {
            version: "2012-10-17".to_string(),
            statement: Vec::new(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PolicyStatement {
    #[serde(rename = "Sid")]
    pub sid: Option<String>,
    #[serde(rename = "Effect")]
    pub effect: PolicyEffect,
    #[serde(rename = "Principal")]
    pub principal: Principal,
    #[serde(rename = "Action")]
    pub action: Vec<String>,
    #[serde(rename = "Resource")]
    pub resource: Vec<String>,
    #[serde(rename = "Condition")]
    pub condition: Option<HashMap<String, HashMap<String, String>>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum PolicyEffect {
    Allow,
    Deny,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Principal {
    Wildcard(String),
    Specific(HashMap<String, Vec<String>>),
}

impl Principal {
    pub fn is_public(&self) -> bool {
        match self {
            Principal::Wildcard(s) => s == "*",
            Principal::Specific(_) => false,
        }
    }

    pub fn matches_user(&self, user_id: &str) -> bool {
        match self {
            Principal::Wildcard(s) => s == "*",
            Principal::Specific(map) => {
                if let Some(aws_users) = map.get("AWS") {
                    aws_users.iter().any(|u| u == user_id || u == "*")
                } else {
                    false
                }
            }
        }
    }
}

impl BucketPolicy {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn is_allowed(&self, action: &str, resource: &str, user_id: &str) -> bool {
        let mut allowed = false;
        
        for stmt in &self.statement {
            if stmt.matches_action(action) && stmt.matches_resource(resource) {
                if stmt.principal.matches_user(user_id) {
                    match stmt.effect {
                        PolicyEffect::Allow => {
                            if stmt.matches_condition(user_id) {
                                allowed = true;
                            }
                        }
                        PolicyEffect::Deny => {
                            if stmt.matches_condition(user_id) {
                                return false;
                            }
                        }
                    }
                }
            }
        }
        
        allowed
    }
}

impl PolicyStatement {
    pub fn matches_action(&self, action: &str) -> bool {
        self.action.iter().any(|a| {
            a == action || a == "s3:*" || a == "*" ||
            (a.ends_with('*') && action.starts_with(&a[..a.len()-1]))
        })
    }

    pub fn matches_resource(&self, resource: &str) -> bool {
        self.resource.iter().any(|r| {
            r == resource || r == "*" ||
            (r.ends_with('*') && resource.starts_with(&r[..r.len()-1]))
        })
    }

    pub fn matches_condition(&self, _user_id: &str) -> bool {
        if let Some(cond) = &self.condition {
            for (operator, values) in cond {
                for (key, value) in values {
                    if operator == "StringEquals" && key == "aws:userid" {
                        if value != _user_id {
                            return false;
                        }
                    }
                }
            }
        }
        true
    }
}

pub fn default_public_policy(bucket: &str) -> BucketPolicy {
    BucketPolicy {
        version: "2012-10-17".to_string(),
        statement: vec![
            PolicyStatement {
                sid: Some("PublicRead".to_string()),
                effect: PolicyEffect::Allow,
                principal: Principal::Wildcard("*".to_string()),
                action: vec!["s3:GetObject".to_string()],
                resource: vec![format!("arn:aws:s3:::{}/*", bucket)],
                condition: None,
            },
        ],
    }
}

pub fn default_private_policy(bucket: &str, user_id: &str) -> BucketPolicy {
    BucketPolicy {
        version: "2012-10-17".to_string(),
        statement: vec![
            PolicyStatement {
                sid: Some("OwnerFullAccess".to_string()),
                effect: PolicyEffect::Allow,
                principal: Principal::Specific({
                    let mut map = HashMap::new();
                    map.insert("AWS".to_string(), vec![format!("arn:aws:iam:::user/{}", user_id)]);
                    map
                }),
                action: vec!["s3:*".to_string()],
                resource: vec![
                    format!("arn:aws:s3:::{}/*", bucket),
                    format!("arn:aws:s3:::{}/*", bucket),
                ],
                condition: None,
            },
        ],
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_policy_parse() {
        let policy_json = r#"{
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "allow",
                    "Principal": "*",
                    "Action": ["s3:GetObject"],
                    "Resource": ["arn:aws:s3:::mybucket/*"]
                }
            ]
        }"#;

        let policy: BucketPolicy = serde_json::from_str(policy_json).unwrap();
        assert_eq!(policy.version, "2012-10-17");
        assert_eq!(policy.statement.len(), 1);
        assert_eq!(policy.statement[0].effect, PolicyEffect::Allow);
    }

    #[test]
    fn test_policy_evaluation_allow() {
        let policy = default_public_policy("testbucket");
        assert!(policy.is_allowed("s3:GetObject", "arn:aws:s3:::testbucket/file.txt", "anonymous"));
    }

    #[test]
    fn test_policy_evaluation_deny() {
        let policy = default_public_policy("testbucket");
        assert!(!policy.is_allowed("s3:PutObject", "arn:aws:s3:::testbucket/file.txt", "anonymous"));
    }

    #[test]
    fn test_action_wildcard() {
        let stmt = PolicyStatement {
            sid: None,
            effect: PolicyEffect::Allow,
            principal: Principal::Wildcard("*".to_string()),
            action: vec!["s3:*".to_string()],
            resource: vec!["*".to_string()],
            condition: None,
        };

        assert!(stmt.matches_action("s3:GetObject"));
        assert!(stmt.matches_action("s3:PutObject"));
        assert!(stmt.matches_action("s3:DeleteObject"));
    }

    #[test]
    fn test_resource_pattern() {
        let stmt = PolicyStatement {
            sid: None,
            effect: PolicyEffect::Allow,
            principal: Principal::Wildcard("*".to_string()),
            action: vec!["s3:GetObject".to_string()],
            resource: vec!["arn:aws:s3:::mybucket/home/*".to_string()],
            condition: None,
        };

        assert!(stmt.matches_resource("arn:aws:s3:::mybucket/home/user/file.txt"));
        assert!(!stmt.matches_resource("arn:aws:s3:::mybucket/public/file.txt"));
    }

    #[test]
    fn test_principal_user_match() {
        let principal = Principal::Specific({
            let mut map = HashMap::new();
            map.insert("AWS".to_string(), vec!["warren".to_string()]);
            map
        });

        assert!(principal.matches_user("warren"));
        assert!(!principal.matches_user("demo"));
    }
}