use std::borrow::Cow;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::atomic::AtomicI64;
use usiem::components::common::*;
use usiem::crossbeam_channel::{Receiver, Sender, bounded};
use usiem::prelude::alert::SiemAlert;
use usiem::prelude::dataset::text_map::TextMapSynDataset;
use usiem::prelude::dataset::{SiemDatasetType};
use usiem::prelude::dataset::holder::DatasetHolder;
use usiem::prelude::metrics::{SiemMetric, SiemMetricDefinition};
use usiem::prelude::SiemComponent;
use usiem::prelude::*;
use crate::client::{self, NotionClient, NotionResult};
const NOTION_API_KEY : &'static str = "API_KEY";
const NOTION_DATABASE : &'static str = "DATABASE_ID";
#[derive(Clone)]
struct NotionMetrics {
pub generated_alerts: SiemMetric,
}
#[derive(Clone)]
pub struct NotionAlert {
id: u64,
local_channel: (Sender<SiemMessage>, Receiver<SiemMessage>),
kernel: Sender<SiemMessage>,
conn: Box<dyn SiemComponentStateStorage>,
datasets: DatasetHolder,
metrics: NotionMetrics,
}
impl NotionAlert {
pub fn new() -> Self {
let local_channel = bounded(1);
let (kernel ,_) = bounded(1);
Self {
local_channel,
id : 0,
datasets : DatasetHolder::new(),
metrics : NotionMetrics {
generated_alerts : SiemMetric::Counter(Arc::new(AtomicI64::new(0)))
},
conn : Box::new(DummyStateStorage{}),
kernel
}
}
}
impl SiemComponent for NotionAlert {
fn set_id(&mut self, id: u64) {
self.id = id;
}
fn local_channel(&self) -> Sender<SiemMessage> {
self.local_channel.0.clone()
}
fn set_log_channel(&mut self, _sender: Sender<SiemLog>, _receiver: Receiver<SiemLog>) {
}
fn set_kernel_sender(&mut self, sender: Sender<SiemMessage>) {
self.kernel = sender;
}
fn run(&mut self) {
let secret_dataset : &TextMapSynDataset = match self.datasets.get(&SiemDatasetType::Secrets(Cow::Borrowed("NotionAlerter"))) {
Some(v) => match v.try_into() {
Ok(v) => v,
Err(_) => return
},
None => return
};
let api_key = match secret_dataset.get(NOTION_API_KEY) {
Some(v) => v,
None => return
};
let database_id = match secret_dataset.get(NOTION_DATABASE) {
Some(v) => v,
None => return
};
let client = client::NotionClient::new(api_key, database_id);
if !client.check_valid_siem_database().unwrap() {
return;
}
loop {
let msg = self.local_channel.1.recv().unwrap();
match msg {
SiemMessage::Alert(alert) => {
match send_alert_with_retry(&client, &alert) {
Ok(_) => {},
Err(_) => {
self.kernel.send(SiemMessage::Notification(self.id, Cow::Owned(format!("Cannot generate alert: {:?}", alert)))).unwrap();
}
}
},
SiemMessage::Command(_,call) => {
match call {
command::SiemCommandCall::STOP_COMPONENT(_) => return,
_ => {},
}
}
_ => {},
}
}
}
fn set_storage(&mut self, conn: Box<dyn SiemComponentStateStorage>) {
self.conn = conn;
}
fn capabilities(&self) -> SiemComponentCapabilities {
let datasets = vec![DatasetDefinition::new(
dataset::SiemDatasetType::Secrets(Cow::Borrowed("NotionAlerter")),
Cow::Borrowed("Access notion API_KEY and DATABASE_ID"),
UserRole::Engineer,
)];
let metrics = vec![SiemMetricDefinition {
metric: self.metrics.generated_alerts.clone(),
name: Cow::Borrowed("generated_alerts"),
description: Cow::Borrowed("Number of alerts generated by this component"),
tags: BTreeMap::new(),
}];
SiemComponentCapabilities::new(
Cow::Borrowed("NotionAlerter"),
Cow::Borrowed("Send alerts to Notion"),
Cow::Borrowed(""),
datasets,
vec![],
vec![],
metrics,
)
}
fn duplicate(&self) -> Box<dyn SiemComponent> {
Box::new(self.clone())
}
fn set_datasets(&mut self, datasets: DatasetHolder) {
self.datasets = datasets;
}
fn id(&self) -> u64 {
return self.id;
}
fn name(&self) -> &str {
return &"NotionAlerter";
}
}
fn send_alert_with_retry(client :&NotionClient, alert :&SiemAlert) -> NotionResult<()>{
let mut errors = 0;
loop {
if let Err(e) = client.send_alert(alert) {
errors += 1;
if errors > 5 {
return Err(e)
}
}else{
return Ok(())
}
}
}
#[cfg(test)]
mod component {
use std::{borrow::Cow, sync::Arc, time::UNIX_EPOCH};
use usiem::{prelude::{dataset::{SiemDataset, text_map::{TextMapSynDataset, TextMapDataset}, holder::DatasetHolder}, SiemComponent, command::{SiemCommandHeader, SiemCommandCall}, mitre::MitreTechniques, alert::{AlertSeverity, SiemAlert}, SiemLog, SiemEvent, auth::{AuthEvent, LoginOutcome, AuthLoginType, RemoteLogin}}, crossbeam_channel::bounded, components::common::SiemMessage};
use crate::NotionAlert;
use super::{NOTION_DATABASE, NOTION_API_KEY};
#[test]
fn shoul_generate_alert() {
let db_id : String = match std::env::var("USIEM_NOTION_DB") {
Ok(v) => v,
Err(_) => return
};
let api_key = std::env::var("USIEM_NOTION_APIKEY").expect("USIEM_NOTION_DB is defined but not USIEM_NOTION_APIKEY");
let mut dataset = TextMapDataset::new();
dataset.insert(Cow::Borrowed(NOTION_DATABASE), Cow::Owned(db_id));
dataset.insert(Cow::Borrowed(NOTION_API_KEY), Cow::Owned(api_key));
let (comm, _recv) = bounded(10);
let dataset = Arc::new(dataset);
let secret_dataset = SiemDataset::Secrets((Cow::Borrowed("NotionAlerter"), TextMapSynDataset::new(dataset, comm)));
let datasets = DatasetHolder::from_datasets(vec![secret_dataset]);
let mut comp = NotionAlert::new();
comp.set_datasets(datasets);
let local_channel = comp.local_channel();
let join = std::thread::spawn(move || {
comp.run();
});
let mut log = SiemLog::new(String::from("2021/02/05 15:05:55,0123456789,TRAFFIC,drop,2305,2021/02/05 15:05:55,192.168.2.1,192.168.3.2,0.0.0.0,0.0.0.0,intrazone-default,,,not-applicable,vsys1,untrust,untrust,ethernet1,,Log Collector,2021/02/05 15:05:55,0,1,35861,514,0,0,0x0,udp,deny,102,102,0,1,2021/02/08 16:06:59,0,any,0,3996658540,0x0,192.168.0.0-192.168.255.255,192.168.0.0-192.168.255.255,0,1,0,policy-deny,0,0,0,0,,pa-test,from-policy,,,0,,0,,N/A,0,0,0,0,aa1d1908-68cd-467a-bc90-68a57096ef3b,0,0,,,,,,,"), 0, "localhost");
log.set_tenant(Cow::Borrowed("Contoso"));
log.set_event(SiemEvent::Auth(AuthEvent {
hostname: Cow::Borrowed("hostname1"),
outcome: LoginOutcome::FAIL,
login_type: AuthLoginType::Remote(RemoteLogin {
domain: Cow::Borrowed("CNMS"),
source_address: Cow::Borrowed("10.10.10.10"),
user_name: Cow::Borrowed("cancamusa"),
}),
}));
let alert = SiemAlert {
title: format!("(TEST) shoul_generate_alert {:?}", std::time::SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs()),
description: String::from("This is a test of the NotionAlerter component"),
severity: AlertSeverity::CRITICAL,
date: usiem::chrono::Utc::now().timestamp_millis(),
tags: vec![String::from("Critical")],
techniques : vec![MitreTechniques::T1001, MitreTechniques::T1003_001],
rule: String::from("ruleset::example::rule1"),
log: log.clone(),
aggr_limit: 0,
aggr_key: String::from("example::rule"),
};
local_channel.send(SiemMessage::Alert(alert)).unwrap();
local_channel.send(SiemMessage::Command(SiemCommandHeader{
comm_id :0,
comp_id : 0,
user : format!("zero")
}, SiemCommandCall::STOP_COMPONENT(String::new()))).unwrap();
join.join().unwrap();
}
}