use chrono::{DateTime, Datelike, Timelike, Utc};
use kimberlite_types::{ClearanceLevel, DataClass};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum DeviceType {
Desktop,
Mobile,
Server,
Unknown,
}
pub const MAX_CLEARANCE: u8 = ClearanceLevel::TopSecret.as_u8();
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserAttributes {
pub role: String,
pub department: String,
pub clearance_level: ClearanceLevel,
pub ip_address: Option<String>,
pub device_type: DeviceType,
pub tenant_id: Option<u64>,
}
impl UserAttributes {
pub fn new(role: &str, department: &str, clearance_level: u8) -> Self {
let clearance_level =
ClearanceLevel::try_from(clearance_level).unwrap_or(ClearanceLevel::TopSecret);
Self::with_clearance(role, department, clearance_level)
}
pub fn with_clearance(role: &str, department: &str, clearance_level: ClearanceLevel) -> Self {
Self {
role: role.to_string(),
department: department.to_string(),
clearance_level,
ip_address: None,
device_type: DeviceType::Unknown,
tenant_id: None,
}
}
pub fn with_ip(mut self, ip: &str) -> Self {
self.ip_address = Some(ip.to_string());
self
}
pub fn with_device(mut self, device: DeviceType) -> Self {
self.device_type = device;
self
}
pub fn with_tenant(mut self, tenant_id: u64) -> Self {
self.tenant_id = Some(tenant_id);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceAttributes {
pub data_class: DataClass,
pub owner_tenant: u64,
pub stream_name: String,
pub retention_days: Option<u32>,
pub correction_allowed: bool,
pub legal_hold_active: bool,
pub requested_fields: Option<Vec<String>>,
}
impl ResourceAttributes {
pub fn new(data_class: DataClass, owner_tenant: u64, stream_name: &str) -> Self {
Self {
data_class,
owner_tenant,
stream_name: stream_name.to_string(),
retention_days: None,
correction_allowed: false,
legal_hold_active: false,
requested_fields: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnvironmentAttributes {
pub timestamp: DateTime<Utc>,
pub is_business_hours: bool,
pub source_country: String,
}
impl EnvironmentAttributes {
pub fn from_timestamp(ts: DateTime<Utc>, country: &str) -> Self {
let hour = ts.hour();
let weekday = ts.weekday();
let is_weekday = matches!(
weekday,
chrono::Weekday::Mon
| chrono::Weekday::Tue
| chrono::Weekday::Wed
| chrono::Weekday::Thu
| chrono::Weekday::Fri
);
let is_business_hours = is_weekday && (9..17).contains(&hour);
Self {
timestamp: ts,
is_business_hours,
source_country: country.to_string(),
}
}
pub fn new(timestamp: DateTime<Utc>, is_business_hours: bool, source_country: &str) -> Self {
Self {
timestamp,
is_business_hours,
source_country: source_country.to_string(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::TimeZone;
#[test]
fn test_business_hours_weekday_morning() {
let ts = Utc.with_ymd_and_hms(2025, 1, 8, 10, 0, 0).unwrap();
let env = EnvironmentAttributes::from_timestamp(ts, "US");
assert!(
env.is_business_hours,
"10:00 UTC on Wednesday should be business hours"
);
}
#[test]
fn test_business_hours_weekday_evening() {
let ts = Utc.with_ymd_and_hms(2025, 1, 8, 18, 0, 0).unwrap();
let env = EnvironmentAttributes::from_timestamp(ts, "US");
assert!(
!env.is_business_hours,
"18:00 UTC on Wednesday should not be business hours"
);
}
#[test]
fn test_business_hours_weekend() {
let ts = Utc.with_ymd_and_hms(2025, 1, 11, 10, 0, 0).unwrap();
let env = EnvironmentAttributes::from_timestamp(ts, "US");
assert!(
!env.is_business_hours,
"10:00 UTC on Saturday should not be business hours"
);
}
#[test]
fn test_business_hours_boundary_start() {
let ts = Utc.with_ymd_and_hms(2025, 1, 8, 9, 0, 0).unwrap();
let env = EnvironmentAttributes::from_timestamp(ts, "US");
assert!(
env.is_business_hours,
"09:00 UTC on Wednesday should be business hours"
);
}
#[test]
fn test_business_hours_boundary_end() {
let ts = Utc.with_ymd_and_hms(2025, 1, 8, 17, 0, 0).unwrap();
let env = EnvironmentAttributes::from_timestamp(ts, "US");
assert!(
!env.is_business_hours,
"17:00 UTC on Wednesday should not be business hours (exclusive end)"
);
}
#[test]
fn test_user_attributes_builder() {
let user = UserAttributes::new("admin", "engineering", 3)
.with_ip("192.168.1.1")
.with_device(DeviceType::Desktop)
.with_tenant(42);
assert_eq!(user.role, "admin");
assert_eq!(user.department, "engineering");
assert_eq!(user.clearance_level, ClearanceLevel::TopSecret);
assert_eq!(user.ip_address, Some("192.168.1.1".to_string()));
assert_eq!(user.device_type, DeviceType::Desktop);
assert_eq!(user.tenant_id, Some(42));
}
#[test]
fn test_user_attributes_clearance_saturates_to_max() {
assert_eq!(
UserAttributes::new("admin", "eng", 10).clearance_level,
ClearanceLevel::TopSecret
);
for c in [4u8, 10, 42, 172, 255] {
assert_eq!(
UserAttributes::new("admin", "engineering", c).clearance_level,
ClearanceLevel::TopSecret
);
}
for (c, expected) in [
(0, ClearanceLevel::Public),
(1, ClearanceLevel::Confidential),
(2, ClearanceLevel::Secret),
(3, ClearanceLevel::TopSecret),
] {
assert_eq!(
UserAttributes::new("admin", "engineering", c).clearance_level,
expected
);
}
}
#[test]
fn test_user_attributes_with_clearance_typed() {
let user = UserAttributes::with_clearance("analyst", "engineering", ClearanceLevel::Secret);
assert_eq!(user.clearance_level, ClearanceLevel::Secret);
}
#[test]
fn test_resource_attributes() {
let resource = ResourceAttributes::new(DataClass::PHI, 1, "patient_records");
assert_eq!(resource.data_class, DataClass::PHI);
assert_eq!(resource.owner_tenant, 1);
assert_eq!(resource.stream_name, "patient_records");
}
}