extern crate pbr;
extern crate regex;
use std::fs::File;
use std::io::{Read, Write};
use std::net::{IpAddr, Ipv4Addr, Shutdown, TcpStream};
use std::str::from_utf8;
use std::thread;
use lazy_static::lazy_static;
use pbr::{MultiBar, Pipe, ProgressBar, Units};
use regex::Regex;
lazy_static! {
static ref DCC_SEND_REGEX: Regex =
Regex::new(r#"DCC SEND "?(.*)"? (\d+) (\d+) (\d+)"#).unwrap();
static ref PING_REGEX: Regex = Regex::new(r#"PING :\d+"#).unwrap();
static ref JOIN_REGEX: Regex = Regex::new(r#"JOIN :#.*"#).unwrap();
}
pub struct IRCRequest {
pub server: String,
pub channel: String,
pub nickname: String,
pub bot: String,
pub packages: Vec<String>,
}
struct DCCSend {
filename: String,
ip: IpAddr,
port: String,
file_size: usize,
}
pub fn connect_and_download(request: IRCRequest, on_start: fn(String) -> ()) -> Result<(), String> {
let mut download_handles = Vec::new();
let mut has_joined = false;
let mut multi_bar = MultiBar::new();
let mut stream = log_in(&request).unwrap();
let mut message_buffer = String::new();
while download_handles.len() < request.packages.len() {
let message = read_next_message(&mut stream, &mut message_buffer).unwrap();
if PING_REGEX.is_match(&message) {
let pong = message.replace("PING", "PONG");
stream.write(pong.as_bytes()).unwrap();
if !has_joined {
let channel_join_cmd = format!("JOIN #{}\r\n", request.channel);
stream.write(channel_join_cmd.as_bytes()).unwrap();
has_joined = true;
}
}
if JOIN_REGEX.is_match(&message) {
for package in &request.packages {
let xdcc_send_cmd = format!("PRIVMSG {} :xdcc send #{}\r\n", request.bot, package);
stream.write(xdcc_send_cmd.as_bytes()).unwrap();
}
}
if DCC_SEND_REGEX.is_match(&message) {
let request = parse_dcc_send(&message);
let mut progress_bar = multi_bar.create_bar(request.file_size as u64);
let handle = thread::spawn(move || {
download_file(request, &mut progress_bar, on_start).unwrap();
});
download_handles.push(handle);
}
}
stream
.write("QUIT :my job is done here!\r\n".as_bytes())
.unwrap();
stream.shutdown(Shutdown::Both).unwrap();
multi_bar.listen();
download_handles
.into_iter()
.for_each(|handle| handle.join().unwrap());
Ok(())
}
fn log_in(request: &IRCRequest) -> Result<TcpStream, std::io::Error> {
let mut stream = TcpStream::connect(&request.server)?;
stream.write(format!("NICK {}\r\n", request.nickname).as_bytes())?;
stream.write(format!("USER {} 0 * {}\r\n", request.nickname, request.nickname).as_bytes())?;
Ok(stream)
}
fn read_next_message(
stream: &mut TcpStream,
message_builder: &mut String,
) -> Result<String, std::io::Error> {
let mut buffer = [0; 4];
while !message_builder.contains("\n") {
let count = stream.read(&mut buffer[..])?;
message_builder.push_str(from_utf8(&buffer[..count]).unwrap_or_default());
}
let endline_offset = message_builder.find('\n').unwrap() + 1;
let message = message_builder.get(..endline_offset).unwrap().to_string();
message_builder.replace_range(..endline_offset, "");
Ok(message)
}
fn parse_dcc_send(message: &String) -> DCCSend {
let captures = DCC_SEND_REGEX.captures(&message).unwrap();
let ip_number = captures[2].parse::<u32>().unwrap();
DCCSend {
filename: captures[1].to_string(),
ip: IpAddr::V4(Ipv4Addr::from(ip_number)),
port: captures[3].to_string(),
file_size: captures[4].parse::<usize>().unwrap(),
}
}
fn download_file(
request: DCCSend,
progress_bar: &mut ProgressBar<Pipe>,
on_start: fn(String) -> (),
) -> std::result::Result<(), std::io::Error> {
let filename = request.filename.to_string();
let mut file = File::create(&request.filename)?;
let mut stream = TcpStream::connect(format!("{}:{}", request.ip, request.port))?;
let mut buffer = [0; 4096];
let mut progress: usize = 0;
progress_bar.set_units(Units::Bytes);
progress_bar.message(&format!("{}: ", &request.filename));
on_start(filename);
while progress < request.file_size {
let count = stream.read(&mut buffer[..])?;
file.write(&mut buffer[..count])?;
progress += count;
progress_bar.set(progress as u64);
}
progress_bar.finish();
stream.shutdown(Shutdown::Both)?;
file.flush()?;
Ok(())
}