epistlebox/epistlebox.rs

use std::io::prelude::*;
use std::io;
use std::collections::HashSet;
use std::fs::{File,create_dir_all,OpenOptions};
use std::path::PathBuf;
use std::net::{TcpStream,TcpListener,IpAddr};
use std::env::args;
use std::thread;
use std::process::exit;

#[derive(PartialEq)]
enum PortPrivacy {
	Public,
	Whitelist,
	Private,
}

static BLOCKFILE_NAME:&str = ".noepistles";
static ALLOWFILE_NAME:&str = ".yeepistles";
static LOGFILE_NAME:&str = ".epistleslog";
static FILENAME_PREFIX:&str = "epistle";
static FILENAME_SUFFIX:&str = ".epistle";

fn load_rules(rulefilepath:&str) -> HashSet<IpAddr> {
	let mut rules:HashSet<IpAddr> = HashSet::new();
	match File::open(rulefilepath) {
		Err(why) => match why.kind() {
			io::ErrorKind::NotFound => (),
			_ => {
				eprintln!("-!- failed to open {} to read IP address ruleing rules: {}",rulefilepath,why);
				exit(1);
			},
		},
		Ok(mut file) => {
			let mut rulestring = String::new();
			match file.read_to_string(&mut rulestring) {
				Err(why) => {
					eprintln!("-!- failed to read {} for IP address ruleing rules: {}",rulefilepath,why);
					exit(1);
				},
				Ok(_) => (),
			};
			for line in rulestring.lines() {
				match line.parse::<IpAddr>() {
					Err(_why) => {
						eprintln!("-!- could not parse '{}' in {} as an IP address",line,rulefilepath);
						exit(1);
					},
					Ok(address) => rules.insert(address),
				};
			}
		},
	};
	return rules;
}

fn record(mut stream:TcpStream,filename:&str) {
	let mut bottle:[u8;512] = [0x00;512];
	let mut outfile = match File::create(filename) {
		Err(why) => {
			eprintln!("-!- failed to open file {} for writing: {}",filename,why);
			return;
		},
		Ok(f) => f,
	};
	loop {
		match stream.read(&mut bottle) {
			Err(why) => match why.kind() {
				io::ErrorKind::WouldBlock => (),
				io::ErrorKind::Interrupted => (),
				_ => {
					eprintln!("-!- failed to collect data for writing to {}: {}",filename,why);
					return;
				},
			},
			Ok(n) => match n {
				0 => {
					println!("-i- finished receiving {}",filename);
					return;
				},
				_ => match outfile.write_all(&bottle[..n]) {
					Err(why) => {
						eprintln!("-!- failed to write to {}: {}",filename,why);
						return;
					},
					Ok(_) => (),
				},
			},
		};
	}
}

fn main() {
	
	let argv:Vec<String> = args().collect();
	if argv.len() < 2 {
		eprintln!("usage: epistlebox [port number]...");
		exit(1);
	}

	let mut ports:Vec<(PortPrivacy,u16)> = Vec::new();
	for arg in argv[1..].iter() {
		match arg.trim_end_matches(&['w','p'][..]).parse() {
			Err(_why) => {
				eprintln!("-!- could not parse '{}' as a port number",arg);
			},
			Ok(n) => match arg.chars().last() {
				Some('w') => {
					println!("-i- accepting only connections whitelisted in {} on port {}",ALLOWFILE_NAME,n);
					ports.push((PortPrivacy::Whitelist,n));
				},
				Some('p') => {
					println!("-i- accepting only local connections on port {}",n);
					ports.push((PortPrivacy::Private,n));
				},
				_ => ports.push((PortPrivacy::Public,n)),
			},
		};
	}
	
	let mut listeners:Vec<(PortPrivacy,u16,TcpListener)> = Vec::new();
	while let Some((privacy,port)) = ports.pop() {
		let bindaddress = match privacy {
			PortPrivacy::Private => "127.0.0.1",
			_ => "[::]",
		};
		match TcpListener::bind(&format!("{}:{}",bindaddress,port)) {
			Err(why) => {
				eprintln!("-!- could not bind listener to port {}: {}",port,why);
			},
			Ok(listener) => {
				listener.set_nonblocking(true).expect("cannot set listener to nonblocking");
				listeners.push((privacy,port,listener));
			},
		};
	}

	let blocklist = load_rules(BLOCKFILE_NAME);
	let allowlist = load_rules(ALLOWFILE_NAME);

	loop {
		for (privacy,port,listener) in listeners.iter() {
			match listener.accept() {
				Err(why) => match why.kind() {
					io::ErrorKind::WouldBlock => (),
					_ => eprintln!("-!- failed to accept connection on port {}: {}",port,why),
				},
				Ok((stream,address)) => {
					if blocklist.contains(&address.ip()) {
						eprintln!("-x- connection refused for {} (address found in {})",address,BLOCKFILE_NAME);
						continue;
					}
					if privacy == &PortPrivacy::Whitelist && !allowlist.contains(&address.ip()) {
						eprintln!("-x- connection refused for {} (address not found in {})",address,ALLOWFILE_NAME);
						continue;
					}
					let mut filename = format!("./{}{}",FILENAME_PREFIX,FILENAME_SUFFIX);
					let mut filenumber = 1;
					while PathBuf::from(&filename).exists() {
						filename = format!("{}/{}{}{}",port,FILENAME_PREFIX,filenumber,FILENAME_SUFFIX);
						filenumber += 1;
					}
					println!("-i- receiving {} from {} on port {}...",filename,address,port);
					match create_dir_all(&format!("./{}",port)) {
						Err(why) => {
							eprintln!("-!- failed to access directiory ./{}: {}",port,why);
						},
						Ok(_) => (),
					};
					let threadfilename = filename.clone();
					let logstring = format!("{}: {} via port {}\n",filename,address,port);
					let mut logfile = match OpenOptions::new().append(true).open(&LOGFILE_NAME) {
						Err(why) => match why.kind() {
							io::ErrorKind::NotFound => match File::create(&LOGFILE_NAME) {
								Err(why) => {
									eprintln!("-!- failed to create logfile {}: {}",LOGFILE_NAME,why);
									None
								},
								Ok(file) => Some(file),
							},
							_ => {
								eprintln!("-!- failed to open logfile {}: {}",LOGFILE_NAME,why);
								None
							},
						}
						Ok(file) => Some(file),
					};
					if let Some(mut file) = logfile {
						match file.write_all(&logstring.as_bytes()) {
							Err(why) => eprintln!("-!- failed to write logfile {}: {}",LOGFILE_NAME,why),
							Ok(_) => (),
						};
					}
					thread::spawn(move || {
						record(stream,&threadfilename);
					});
				},
			};
		}
	}
}