Initial: simple IRC bot implementation
This commit is contained in:
286
src/main.rs
Normal file
286
src/main.rs
Normal file
@@ -0,0 +1,286 @@
|
||||
use std::net::TcpStream;
|
||||
use std::io::{self, Read, Write};
|
||||
use std::ops::ControlFlow;
|
||||
|
||||
#[derive(Debug)]
|
||||
#[allow(dead_code)]
|
||||
enum IrcClientMessage<'msg> {
|
||||
Raw(&'msg [u8]),
|
||||
RawLine(&'msg [u8]),
|
||||
Nick(&'msg str),
|
||||
User(&'msg str, &'msg str),
|
||||
Pong(&'msg str),
|
||||
Notice(&'msg str, &'msg str),
|
||||
}
|
||||
|
||||
impl<'msg> From<IrcClientMessage<'msg>> for Vec<u8> {
|
||||
fn from(value: IrcClientMessage) -> Self {
|
||||
match value {
|
||||
IrcClientMessage::Nick(nick) =>
|
||||
format!("NICK {nick}").into_bytes(),
|
||||
IrcClientMessage::User(ident, gecos) =>
|
||||
format!("USER {ident} 0 * :{gecos}").into_bytes(),
|
||||
IrcClientMessage::Pong(s) =>
|
||||
format!("PONG :{s}").into_bytes(),
|
||||
IrcClientMessage::Notice(to, msg) =>
|
||||
format!("NOTICE {to} :{msg}").into_bytes(),
|
||||
IrcClientMessage::RawLine(bytes) =>
|
||||
bytes.into(),
|
||||
IrcClientMessage::Raw(bytes) =>
|
||||
bytes.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
#[allow(dead_code)]
|
||||
enum IrcServerMessage<'msg> {
|
||||
RawLine(&'msg [u8]),
|
||||
Ping(&'msg str),
|
||||
Notice{ from: &'msg str, to: &'msg str, msg: &'msg str },
|
||||
Connected { from: &'msg str, nick: &'msg str, msg: &'msg str },
|
||||
}
|
||||
|
||||
struct IrcProtocol(TcpStream);
|
||||
|
||||
impl<'msg> IrcProtocol {
|
||||
fn get_stream(&mut self) -> &mut TcpStream {
|
||||
&mut self.0
|
||||
}
|
||||
|
||||
fn send_message(&mut self, msg: IrcClientMessage) -> io::Result<()> {
|
||||
let is_raw = matches!(msg, IrcClientMessage::Raw(_));
|
||||
|
||||
// translate message into bytes
|
||||
let mut msg_bytes: Vec<u8> = msg.into();
|
||||
|
||||
// leave Raw messages as they are - without extra newlines
|
||||
if !is_raw {
|
||||
msg_bytes.push(b'\r');
|
||||
msg_bytes.push(b'\n');
|
||||
}
|
||||
|
||||
// print what we're sending
|
||||
let msg_str: String = msg_bytes.iter().map(|&c| c as char).collect();
|
||||
print!("[<<] {msg_str}");
|
||||
|
||||
// send the bytes
|
||||
self.get_stream().write_all(&msg_bytes)
|
||||
}
|
||||
|
||||
fn resolve_message(&self, msg_bytes: &'msg [u8]) -> io::Result<IrcServerMessage<'msg>> {
|
||||
assert!(msg_bytes.iter().all(|&c| c != b'\r' && c != b'\n'), "msg_bytes ({:?}) has newline characters", msg_bytes);
|
||||
|
||||
if msg_bytes.starts_with(b"PING :") {
|
||||
return Ok(IrcServerMessage::Ping(str::from_utf8(&msg_bytes[6..]).unwrap()));
|
||||
}
|
||||
|
||||
if msg_bytes.len() > 0 && msg_bytes[0] == b':' {
|
||||
// the message if likely of the ":<server> ..." or ":nick!user@host ..." format;
|
||||
// those tend to have 3+ space-spearated parts to them
|
||||
let msg_parts: Vec<&[u8]> = msg_bytes
|
||||
.splitn(4, |&c| c == b' ')
|
||||
.collect();
|
||||
|
||||
if msg_parts.len() >= 3 {
|
||||
// convert msg_parts into an str vector
|
||||
let mut parts: Vec<&str> = Vec::new();
|
||||
let parse_result = msg_parts
|
||||
.iter()
|
||||
.try_for_each(|&part|
|
||||
match str::from_utf8(part) {
|
||||
Ok(s) => ControlFlow::Continue(parts.push(s)),
|
||||
Err(e) => ControlFlow::Break((part, e)),
|
||||
}
|
||||
);
|
||||
|
||||
if let ControlFlow::Break((part, e)) = parse_result {
|
||||
let strpart: String = part.iter().map(|&c| c as char).collect();
|
||||
eprintln!("[!!] from_utf8({strpart:?}) failed -> {e:?}");
|
||||
} else {
|
||||
// extract the various str components of the message
|
||||
let from = parts[0];
|
||||
let cmd = parts[1];
|
||||
let to = parts[2];
|
||||
let rest = if parts.len() >= 4 {
|
||||
parts[3]
|
||||
} else {
|
||||
""
|
||||
};
|
||||
|
||||
// if we found a suitable message, return it
|
||||
if let Some(msg) = self._resolve_triple(from, cmd, to, rest) {
|
||||
return Ok(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// couldn't find a suitable message - resort to just returning a
|
||||
// raw one
|
||||
Ok(IrcServerMessage::RawLine(msg_bytes))
|
||||
}
|
||||
|
||||
fn _resolve_triple(
|
||||
&self,
|
||||
from: &'msg str,
|
||||
cmd: &'msg str,
|
||||
to: &'msg str,
|
||||
rest: &'msg str,
|
||||
) -> Option<IrcServerMessage<'msg>> {
|
||||
match cmd {
|
||||
"NOTICE" => Some(IrcServerMessage::Notice { from, to, msg: rest }),
|
||||
"001" => Some(IrcServerMessage::Connected { from, nick: to, msg: rest }),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Client {
|
||||
nick: String,
|
||||
ident: String,
|
||||
gecos: String,
|
||||
_proto: Option<IrcProtocol>,
|
||||
_mynick: Option<String>,
|
||||
_myserver: Option<String>,
|
||||
}
|
||||
|
||||
impl Client {
|
||||
fn new(nick: &str, ident: &str, gecos: &str) -> Self {
|
||||
Self {
|
||||
nick: nick.to_string(),
|
||||
ident: ident.to_string(),
|
||||
gecos: gecos.to_string(),
|
||||
_proto: None,
|
||||
_mynick: None,
|
||||
_myserver: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn protocol(&self) -> &IrcProtocol {
|
||||
self._proto.as_ref().unwrap()
|
||||
}
|
||||
|
||||
fn protocol_mut(&mut self) -> &mut IrcProtocol {
|
||||
self._proto.as_mut().unwrap()
|
||||
}
|
||||
|
||||
fn connect(&mut self, addr: &str) -> io::Result<()> {
|
||||
// make sure we're not already connected
|
||||
if self._proto.is_some() {
|
||||
return Err(io::Error::new(io::ErrorKind::AlreadyExists, "already connected"));
|
||||
}
|
||||
|
||||
// try to connect
|
||||
match TcpStream::connect(addr) {
|
||||
Ok(stream) => {
|
||||
self._proto = Some(IrcProtocol(stream));
|
||||
self.handle_connected()
|
||||
},
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
fn serve(&mut self) -> io::Result<()> {
|
||||
let mut rbuf = [0u8; 4096];
|
||||
let mut filled = 0;
|
||||
|
||||
loop {
|
||||
let sz = self.protocol_mut().get_stream().read(&mut rbuf[filled..])?;
|
||||
let available = filled + sz;
|
||||
if sz == 0 {
|
||||
println!("[--] CONNECTION CLOSED ({filled} bytes lost)");
|
||||
if let Err(e) = self.handle_disconnected() {
|
||||
eprintln!("[EE] handle_disconnected() error: {e:?}");
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
// go over the buffer so far and extract any message that are ready
|
||||
// we know for sure that up to <filled> there are no newline chars
|
||||
let mut start = 0;
|
||||
for i in filled..available {
|
||||
if rbuf[i] == b'\n' {
|
||||
let mut end = i;
|
||||
while end > 0 && rbuf[end-1] == b'\r' {
|
||||
end -= 1;
|
||||
}
|
||||
|
||||
// translate into a message
|
||||
let msg_bytes = &rbuf[start..end];
|
||||
let msg_str: String = msg_bytes.iter().map(|&c| c as char).collect();
|
||||
println!("[>>] {msg_str}");
|
||||
|
||||
let msg = self.protocol().resolve_message(msg_bytes);
|
||||
if msg.is_err() {
|
||||
eprintln!("[EE] Error: {:?} for message {:?}", msg.unwrap_err(), msg_bytes);
|
||||
continue;
|
||||
}
|
||||
|
||||
// handle the message before moving on to the next one
|
||||
if let Err(e) = self.handle_message(msg.unwrap()) {
|
||||
eprintln!("[EE] Error: handle_message() returned {e:?} for {msg_bytes:?}");
|
||||
}
|
||||
start = i+1;
|
||||
}
|
||||
}
|
||||
|
||||
// copy what's left to the beginning of the buffer
|
||||
// TODO: use ring buffer
|
||||
for (dst, i) in (start..available).enumerate() {
|
||||
rbuf[dst] = rbuf[i];
|
||||
}
|
||||
|
||||
filled = available - start;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn send_msg(&mut self, msg: IrcClientMessage) -> io::Result<()> {
|
||||
self.protocol_mut().send_message(msg)
|
||||
}
|
||||
|
||||
fn handle_connected(&mut self) -> io::Result<()> {
|
||||
self.send_msg(IrcClientMessage::Nick(&self.nick.clone()))?;
|
||||
self.send_msg(IrcClientMessage::User(&self.ident.clone(), &self.gecos.clone()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_disconnected(&mut self) -> io::Result<()> {
|
||||
println!("[!!] DISCONNECTED");
|
||||
self._myserver = None;
|
||||
self._mynick = None;
|
||||
self._proto = None;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_message(&mut self, msg: IrcServerMessage) -> io::Result<()> {
|
||||
match msg {
|
||||
IrcServerMessage::Connected { from, nick, msg: _ } => {
|
||||
self._mynick = Some(nick.to_string());
|
||||
self._myserver = Some(from.to_string());
|
||||
self.send_msg(IrcClientMessage::Notice("IceDragon", "Rawr!"))?;
|
||||
},
|
||||
IrcServerMessage::Notice { from, to, msg: _ } => {
|
||||
if to == "AUTH" && self._myserver.is_none() {
|
||||
self._myserver = Some(from.to_string());
|
||||
}
|
||||
},
|
||||
IrcServerMessage::Ping(data) => {
|
||||
self.send_msg(IrcClientMessage::Pong(data))?;
|
||||
},
|
||||
IrcServerMessage::RawLine(_) => {
|
||||
// nothing
|
||||
},
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn main() -> io::Result<()> {
|
||||
let mut client = Client::new("RustBot", "icedragon", "IceDragon's Rust client thingy");
|
||||
client.connect("irc.quickfox.net:6667")?;
|
||||
client.serve()
|
||||
}
|
||||
Reference in New Issue
Block a user