elseware/src/patch/main.rs

#![allow(unused_imports)]

use std::collections::HashMap;
use std::net::{TcpListener, SocketAddr, Ipv4Addr};
use std::net;
use std::thread;
use std::fs;
use std::path::{Path, PathBuf, Components};
use std::convert::AsRef;
use mio::{Events, Poll, Token, Ready, PollOpt};
use rand::{Rng, RngCore};
use libpso::{PacketParseError, PSOPacket};
use libpso::patch::packet::*;
use libpso::crypto::{CipherError, PSOCipher, NullCipher};
use libpso::crypto::pc::PSOPCCipher;
use elseware::common::{send_packet, recv_packet, PacketNetworkError};

const PATCH_PORT: u16 = 11000;
const DATA_PORT: u16 = 11001;


#[derive(Debug)]
enum PatchError {
    PacketNetworkError(PacketNetworkError),
    UnexpectedPacket(Box<dyn PSOPacket>),
    CloseConnection,
}

impl From<PacketNetworkError> for PatchError {
    fn from(err: PacketNetworkError) -> PatchError {
        PatchError::PacketNetworkError(err)
    }
}

#[derive(Debug)]
pub enum PatchPacket {
    PatchWelcomeReply(PatchWelcomeReply),
    LoginReply(LoginReply),
}

impl PatchPacket {
    pub fn from_bytes(data: &Vec<u8>) -> Result<PatchPacket, PacketParseError> {
        match data[2] {
            0x02 => Ok(PatchPacket::PatchWelcomeReply(PatchWelcomeReply::from_bytes(data)?)),
            0x04 => Ok(PatchPacket::LoginReply(LoginReply::from_bytes(data)?)),
            _ => Err(PacketParseError::WrongPacketForServerType)
        }
    }
}

struct Client {
    running: bool,
    key_in: u32,
    key_out: u32,
    socket: mio::tcp::TcpStream,
    cipher_in: Box<dyn PSOCipher>,
    cipher_out: Box<dyn PSOCipher>,
    patch_file_tree: PatchFileTree,
    patch_file_lookup: HashMap<u32, PathBuf>,
}

impl Client {
    fn new(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PathBuf>) -> Client {
        let mut rng = rand::thread_rng();
        let key_in: u32 = rng.gen();
        let key_out: u32 = rng.gen();

        Client {
            running: true,
            key_in: key_in,
            key_out: key_out,
            socket: mio::tcp::TcpStream::from_stream(socket).unwrap(),
            cipher_in: Box::new(NullCipher {}),
            cipher_out: Box::new(NullCipher {}),
            patch_file_tree: patch_file_tree,
            patch_file_lookup: patch_file_lookup,
        }
    }

    fn send(&mut self, pkt: &dyn PSOPacket) {
        match send_packet(&mut self.socket, &mut *self.cipher_out, pkt) {
            Ok(_) => {
                println!("[patch] send ({:?}): {:?}", self.socket, pkt);
            },
            Err(err) => {
                println!("[patch] error sending packet to {:?}: {:?}", self.socket, err);
                self.running = false;
            }
        }
    }

}

#[derive(Debug, Clone)]
enum PatchFileTree {
    Directory(PathBuf, Vec<PatchFileTree>),
    File(PathBuf, u32), // file_id
}

fn load_patch_dir(basedir: &str, patchbase: &str, file_ids: &mut HashMap<u32, PathBuf>) -> PatchFileTree {
    let paths = fs::read_dir(basedir).unwrap();

    let mut files = Vec::new();
    let mut dirs = Vec::new();
    for p in paths {
        let path = p.unwrap().path();
        if path.is_dir() {
            let patch_path = path.strip_prefix(basedir).unwrap();
            dirs.push(load_patch_dir(path.to_str().unwrap(), patch_path.to_str().unwrap(), file_ids));
        }
        else {
            let patch_path = path.strip_prefix(basedir).unwrap();
            files.push(PatchFileTree::File(patch_path.to_path_buf(), file_ids.len() as u32));
            file_ids.insert(file_ids.len() as u32, path);
        }
    }

    files.append(&mut dirs);

    PatchFileTree::Directory(PathBuf::from(patchbase), files)
}

fn generate_patch_tree(basedir: &str) -> (PatchFileTree, HashMap<u32, PathBuf>) {
    let mut file_ids = HashMap::new();

    let patch_tree = load_patch_dir(basedir, "", &mut file_ids);
    
    (patch_tree, file_ids)
}


fn get_file_list_packets(patch_file_tree: &PatchFileTree) -> Vec<Box<dyn PSOPacket>> {
    let mut pkts: Vec<Box<dyn PSOPacket>> = Vec::new();

    match patch_file_tree {
        PatchFileTree::Directory(dir, files) => {
            pkts.push(Box::new(ChangeDirectory::new(dir.to_str().unwrap())));
            for file in files {
                pkts.append(&mut get_file_list_packets(&file));
            }
            pkts.push(Box::new(UpOneDirectory {}));
        },
        PatchFileTree::File(path, id) => {
            pkts.push(Box::new(FileInfo::new(path.to_str().unwrap(), *id)));
        }
    }

    pkts
}


// so I can not figure out why this doesnt work!
fn send_file_list(client: &mut Client) -> Result<(), PatchError> {
    client.send(&PatchStartList {});

    let pkts = get_file_list_packets(&client.patch_file_tree);

    /*for pkt in pkts {
        client.send(&*pkt);
    }*/
    
    client.send(&ChangeDirectory::new(""));
    client.send(&PatchEndList {});
    client.send(&EndIt {});

    Ok(())
}


fn handle_packet(client: &mut Client, pkt: PatchPacket) -> Result<(), PatchError> {
    println!("[patch] recv({:?}): {:?}", client.socket, pkt);

    match pkt {
        PatchPacket::PatchWelcomeReply(pkt) => {
            client.send(&RequestLogin {});
        },
        PatchPacket::LoginReply(pkt) => {
            client.send(&Message::new("hello player".to_string()));
            send_file_list(client);
        },
    }

    Ok(())
}

fn client_loop(mut client: Client) {
    let poll = mio::Poll::new().unwrap();
    poll.register(&client.socket, Token(0), Ready::readable(), PollOpt::edge());

    let mut events = Events::with_capacity(1024);
    loop {
        println!("about to poll {:?}", client.socket);
        poll.poll(&mut events, None).unwrap();
        println!("polled!");
        
        for event in &events{
            println!("event! {:?}", event);
            if event.token() == Token(0) {
                let pkt = recv_packet(&mut client.socket, &mut *client.cipher_in)
                    .and_then(|pkt| {
                        PatchPacket::from_bytes(&pkt)
                            .map_err(|err| err.into())
                    });

                match pkt {
                    Ok(pkt) => {
                        handle_packet(&mut client, pkt);
                    },
                    Err(err) => {
                        println!("[patch] error recv-ing packet with {:?}: {:?}", client.socket, err);
                    }
                }
            }
        }
    }
}

fn new_client(socket: net::TcpStream, patch_file_tree: PatchFileTree, patch_file_lookup: HashMap<u32, PathBuf>) {
    let mut client = Client::new(socket, patch_file_tree, patch_file_lookup);
    let welcome_pkt = PatchWelcome::new(client.key_out, client.key_in);
    client.send(&welcome_pkt);
    client.cipher_in = Box::new(PSOPCCipher::new(client.key_in));
    client.cipher_out = Box::new(PSOPCCipher::new(client.key_out));
    client_loop(client);
}


fn main() {
    println!("[patch] starting server");

    let listener = TcpListener::bind(&SocketAddr::from((Ipv4Addr::new(0,0,0,0), PATCH_PORT))).unwrap();

    let (patch_file_tree, patch_file_lookup) = generate_patch_tree("patchfiles/");
    
    while let Ok((socket, addr)) = listener.accept() {
        let local_patch_file_tree = patch_file_tree.clone();
        let local_patch_file_lookup = patch_file_lookup.clone();
        thread::spawn(move || {
            println!("[patch] accepted connection: {}", addr);
            new_client(socket, local_patch_file_tree, local_patch_file_lookup);
        });
    }

    
    println!("[patch] exiting...");
}