use lettre::{
Transport,
message::{Mailbox, MessageBuilder},
transport::smtp::{SmtpTransport, authentication::Credentials, client::Tls},
};
use reqwest::Client;
use tracing::{error, info};
use crate::error::{Result, TaskFlowError};
use crate::task::{Task, TaskHandler, TaskResult};
pub struct NotificationTaskHandler {
client: Client,
}
impl NotificationTaskHandler {
pub fn new() -> Self {
Self {
client: Client::new(),
}
}
}
#[async_trait::async_trait]
impl TaskHandler for NotificationTaskHandler {
fn task_type(&self) -> &'static str {
"notification"
}
async fn execute(&self, task: &Task) -> Result<TaskResult> {
let notification_type = task
.definition
.payload
.get("type")
.and_then(|v| v.as_str())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing notification type".to_string())
})?;
info!(
"Sending {} notification: {}",
notification_type, task.definition.id
);
match notification_type {
"webhook" => self.send_webhook(task).await,
"email" => self.send_email(task).await,
"slack" => self.send_slack(task).await,
"discord" => self.send_discord(task).await,
nt => Err(TaskFlowError::InvalidConfiguration(format!(
"Unknown notification type: {}",
nt
))),
}
}
}
impl NotificationTaskHandler {
async fn send_webhook(&self, task: &Task) -> Result<TaskResult> {
let url = task
.definition
.payload
.get("url")
.and_then(|v| v.as_str())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing webhook URL".to_string())
})?;
let payload = task.definition.payload.get("payload").cloned();
let response = self
.client
.post(url)
.json(&payload.unwrap_or(serde_json::Value::Object(Default::default())))
.send()
.await
.map_err(|e| TaskFlowError::ExecutionError(e.to_string()))?;
if !response.status().is_success() {
return Err(TaskFlowError::ExecutionError(format!(
"Webhook failed with status: {}",
response.status()
)));
}
Ok(TaskResult {
success: true,
output: Some("Webhook sent successfully".to_string()),
error: None,
execution_time_ms: 0,
metadata: Default::default(),
})
}
async fn send_email(&self, task: &Task) -> Result<TaskResult> {
let smtp_server = task
.definition
.payload
.get("smtp_server")
.and_then(|v| v.as_str())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing SMTP server".to_string())
})?;
let from_email = task
.definition
.payload
.get("from")
.and_then(|v| v.as_str())
.ok_or_else(|| TaskFlowError::InvalidConfiguration("Missing from email".to_string()))?;
let to_email = task
.definition
.payload
.get("to")
.and_then(|v| v.as_str())
.ok_or_else(|| TaskFlowError::InvalidConfiguration("Missing to email".to_string()))?;
let subject = task
.definition
.payload
.get("subject")
.and_then(|v| v.as_str())
.unwrap_or("TaskFlow Notification");
let body = task
.definition
.payload
.get("body")
.and_then(|v| v.as_str())
.unwrap_or("Task completed successfully");
info!(
"Sending email notification for task: {}",
task.definition.id
);
let email = MessageBuilder::new()
.from(
from_email
.parse::<Mailbox>()
.map_err(|e| TaskFlowError::InvalidConfiguration(e.to_string()))?,
)
.to(to_email
.parse::<Mailbox>()
.map_err(|e| TaskFlowError::InvalidConfiguration(e.to_string()))?)
.subject(subject)
.body(body.to_string())
.map_err(|e| TaskFlowError::ExecutionError(e.to_string()))?;
let mut mailer_builder = SmtpTransport::relay(smtp_server)
.map_err(|e| TaskFlowError::ExecutionError(e.to_string()))?;
if let (Some(username), Some(password)) = (
task.definition
.payload
.get("username")
.and_then(|v| v.as_str()),
task.definition
.payload
.get("password")
.and_then(|v| v.as_str()),
) {
mailer_builder = mailer_builder
.credentials(Credentials::new(username.to_string(), password.to_string()));
}
let tls_mode = task
.definition
.payload
.get("tls")
.and_then(|v| v.as_str())
.unwrap_or("required");
let mailer = match tls_mode {
"required" => mailer_builder.build(),
"none" => mailer_builder.tls(Tls::None).build(),
_ => mailer_builder.build(),
};
match mailer.send(&email) {
Ok(_) => {
info!("Email sent successfully to: {}", to_email);
Ok(TaskResult {
success: true,
output: Some(format!("Email sent successfully to {}", to_email)),
error: None,
execution_time_ms: 0,
metadata: Default::default(),
})
}
Err(e) => {
error!("Failed to send email: {}", e);
Ok(TaskResult {
success: false,
output: None,
error: Some(format!("Email sending failed: {}", e)),
execution_time_ms: 0,
metadata: Default::default(),
})
}
}
}
async fn send_slack(&self, task: &Task) -> Result<TaskResult> {
let webhook_url = task
.definition
.payload
.get("webhook_url")
.and_then(|v| v.as_str())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing Slack webhook URL".to_string())
})?;
let message = task
.definition
.payload
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("Task completed");
let payload = serde_json::json!({
"text": message,
"attachments": [{
"color": "#36a64f",
"fields": [
{
"title": "Task ID",
"value": task.definition.id,
"short": true
},
{
"title": "Task Name",
"value": task.definition.name,
"short": true
}
]
}]
});
let response = self
.client
.post(webhook_url)
.json(&payload)
.send()
.await
.map_err(|e| TaskFlowError::ExecutionError(e.to_string()))?;
if !response.status().is_success() {
return Err(TaskFlowError::ExecutionError(format!(
"Slack notification failed: {}",
response.status()
)));
}
Ok(TaskResult {
success: true,
output: Some("Slack notification sent".to_string()),
error: None,
execution_time_ms: 0,
metadata: Default::default(),
})
}
async fn send_discord(&self, task: &Task) -> Result<TaskResult> {
let webhook_url = task
.definition
.payload
.get("webhook_url")
.and_then(|v| v.as_str())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing Discord webhook URL".to_string())
})?;
let message = task
.definition
.payload
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("Task completed");
let payload = serde_json::json!({
"content": message,
"embeds": [{
"title": task.definition.name,
"description": format!("Task ID: {}", task.definition.id),
"color": 5814783,
"timestamp": chrono::Utc::now().to_rfc3339()
}]
});
let response = self
.client
.post(webhook_url)
.json(&payload)
.send()
.await
.map_err(|e| TaskFlowError::ExecutionError(e.to_string()))?;
if !response.status().is_success() {
return Err(TaskFlowError::ExecutionError(format!(
"Discord notification failed: {}",
response.status()
)));
}
Ok(TaskResult {
success: true,
output: Some("Discord notification sent".to_string()),
error: None,
execution_time_ms: 0,
metadata: Default::default(),
})
}
}