Mention2Mail/src/servers.rs

/* Copyright 2021 Daniel Mowitz
 * This file is part of Mention2Mail.
 *
 * Mention2Mail is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Mention2Mail is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with Mention2Mail.  If not, see <https://www.gnu.org/licenses/>.
 */ 

use std::{
	collections::VecDeque,
	collections::HashMap,
	sync::mpsc::Sender,
};
use irc_proto::command::{
	CapSubCommand,
	Command,
};
use toml::map::{
	Map,
	Keys
};
use irc_proto::message::Message;
use toml::value::Value;

/// A static implementation of a ringbuffer.
/// It holds arrays of three strings each,
/// which are meant to represent the author and message respectively.
pub struct MessageRingBuf<T> {
	buf: Vec<Option<T>>,
	ptr: usize,
}

impl <T: Clone> MessageRingBuf<T> {
	pub fn new(size: usize) -> Self {
		
		Self{
			buf: vec!(None; size),
			ptr: 0
		}
	}
	
	/// Adds a message to the end of the buffer,
	/// overwriting the first value if the buffer is full.
	pub fn push(&mut self, val: T) {
		self.ptr = (self.ptr + 1) % self.buf.len();
		self.buf[self.ptr] = Some(val);
	}

	/// Returns the last element of the buffer,
	/// replacing it with none.
	pub fn pop(&mut self) -> Option<T> {
		let l = self.buf.len();
		self.ptr = (self.ptr + l - 1) % l;
		self.buf[self.ptr].take() 
	}

	/// Returns the contained buffer in the correct order.
	pub fn get_queue(& self) -> Vec<Option<T>> {
		let l = self.buf.len();
		let mut queue = self.buf.clone();
		
		for i in self.ptr + 1 .. l {
			queue[i - self.ptr - 1] = self.buf[i].clone();
		}

		for i in 0 .. self.ptr + 1{
			queue[l - self.ptr - 1 + i] = self.buf[i].clone();
		}

		return queue
	}
}

/// Constructs a VecDeque with the messages
/// required to identify with an IRC server.
/// Uses the credentials set in -`config.toml`.
fn get_irc_identify_messages (config: &Value) 
	-> Result<VecDeque<Message>, Box<dyn std::error::Error>> {	
	let mut queue = VecDeque::new();
	
	queue.push_back(Message::from(
		Command::CAP(
			None, 
			CapSubCommand::END, 
			None, 
			None
		)));
   
	match config.get("password") {
		Some(p) => queue.push_back(Message::from(
				Command::PASS(
					String::from(p.as_str().ok_or("Could not parse password.")?)
				)
		)),
		None => ()
	}

	let nick: String;
	match config.get("nickname") {
		Some(n) => {
			nick = String::from(n.as_str().ok_or("Could not parse nickname.")?);
			queue.push_back(Message::from(Command::NICK(nick.clone())));
		},
		None => return Err("No nickname supplied!".into()),
	}

	match config.get("username") {
		Some(u) => queue.push_back(Message::from(
			Command::USER(
				String::from(u.as_str().ok_or("Could not parse username.")?),
				"0".to_owned(),
				String::from(u.as_str().ok_or("Could not parse username.")?)
			)
		)),
		None => queue.push_back(Message::from(
				// nick.clone() is only used once because the value
				// can be moved the second time
				Command::USER(nick.clone(), "0".to_owned(), nick)
		))
	}
	Ok(queue)
}

/// Appends a given VecDeque to include the messages
/// used to join the channels defined in `config.toml`.
fn get_irc_join_messages(channels: &Vec<Value>, queue: VecDeque<Message>)
	-> Result<VecDeque<Message>, Box<dyn std::error::Error>> {
	let mut queue = queue;

	for channel in channels {
		queue.push_back(Message::from(
			Command::JOIN(
				String::from(channel.as_str().ok_or("Could not parse one of the channels")?),
				None,
				None
			)
		))
	}

	Ok(queue)
}

/// Checks messages for mentions of users or commands in private messages.
/// Returns either one of:
/// - None and None in the case no mention or command is recognized
/// - The appropriate answer as Message and None for a command 
/// - None and the mentioned user and channel as String and the context as a MessageRingBuf for mentions.
fn handle_message (message: Message, nick: &String, 
	client_names: Keys, contexts: &mut HashMap<String, MessageRingBuf<[String; 2]>>) 
	-> (Option<Message>, Option<(String, String, Vec<Option<[String; 2]>>)>) {
	
	let sender = match message.clone().source_nickname() {
		Some(s) => String::from(s),
		None => String::from("anonymous")
	};
	match message.clone().command {
		Command::PING(ref data, _) => {
			return (
				Some(Message::from(
					Command::PONG(data.to_owned(), None)
				)), 
				None
			);
		},
		Command::PRIVMSG(ref rec, ref msg) => {
			// Send bot info on receiving a private message.
			if rec.as_str() == nick.as_str() {
				return (
					Some(Message::from(
							Command::PRIVMSG(
								sender.clone(),
								super::texts::get_bot_info(),
							)
						)
					),
					None
				)
			} else {
				match contexts.get_mut(&rec.clone()) {
					Some(ref mut buf) => {
						buf.push([sender.clone(), msg.clone()]);
						for c_name in client_names {
							if msg.contains(c_name) {
								return (None, Some((
									c_name.clone(),
									rec.clone(),
									buf.get_queue(),	
								)))
							}
						}
					}
					None => {
						println!("Channel not recognized: {}", rec.clone());
					}
				}
			}
		},
		_ => println!("{}", message.clone().to_string())
	}
	return (None, None)
}

/// Connects to a server and indefinitely iterates the messages received.
/// Calls handle_message on each one and pushes 
/// email, server, channel as Strings 
/// and the context as MessageRingBuf to the given stream.
pub fn handle_server  
	(config: Value, tx: Sender<(String, String, String, Vec<Option<[String; 2]>>)>,context_size: usize) 
	-> Result<(), Box<dyn std::error::Error>> {
	
	let server_name = config.get("server").unwrap().as_str()
		.ok_or("Could not get server adress from config")?;
	let nick = config.get("nickname").unwrap().as_str()
		.ok_or("Could not get nickname from config")?.to_owned();
	let channels = config.get("channels").unwrap().as_array()
		.ok_or("Could not get channels from config")?;

	let clients;
	match config.get("clients").unwrap().as_table() {
		Some(t) => clients = t.clone(),
		None => clients = Map::new(),
	}
	
	let mut chat_contexts = HashMap::new();
	for channel in channels {
		chat_contexts.insert(
			channel.as_str()
				.ok_or("Could not parse one of the channels")?
				.to_owned(),
			MessageRingBuf::new(context_size),
		);
	}

	let mut stream = super::connect::connect_irc(&config).unwrap();
	let mut message_queue = get_irc_identify_messages(&config).unwrap();

	while let Some(message) = message_queue.pop_front() {
		stream.write(message).unwrap();
	}

	message_queue = get_irc_join_messages(&channels, message_queue).unwrap();

	// Wait for first ping and join channels after sending pong.
	loop {
		let message = stream.read()?;
		match message.command {
			Command::PING(ref data, _) => {
				stream.write(Message::from(
					Command::PONG(data.to_owned(), None)
				)).unwrap();
				while let Some(message) = message_queue.pop_front() {
					stream.write(message).unwrap();
				}
				break;
			}
			_ => ()
		}
	}

	// Handle all incoming messages after joining the channels.
	loop {
		let message = stream.read()?;
		let (answer, data) = handle_message(
			message, 
			&nick, 
			clients.keys(), 
			&mut chat_contexts
		);

		match answer {
			Some(a) => stream.write(a).unwrap(),
			None => ()
		}
		match data {
			Some((client, channel, context)) => match clients.get(&client) {
					Some(addr) => {
						let out = (
						addr.as_str()
							.ok_or("Could not parse email address.")?
							.to_owned(),
						server_name.to_owned(), 
						channel.clone(), 
						context.clone(),
						);
						tx.send(out)?
					},
					None => (),
				},
			None => (),
		}
	}
}