Discover Scope Transport Protocol

tcp-stp-handshake.rs

extern crate encoding;

use std::net::{TcpListener, TcpStream};
use std::str::FromStr;
use std::io::{Read, Write, Result, Error, ErrorKind, BufRead, BufReader};
use encoding::{Encoding, EncoderTrap, DecoderTrap};
use encoding::all::{UTF_16BE, ASCII};

enum STP0Command {
    Services,
    Enable,
    Disable,
    Quit,
}

impl STP0Command {
    fn as_str(&self) -> &'static str {
        match self {
            &STP0Command::Services => "services",
            &STP0Command::Enable => "enable",
            &STP0Command::Disable => "disable",
            &STP0Command::Quit => "quit"
        }
    }
}

struct ScopeConnection {
    socket: TcpStream
}

impl ScopeConnection {
    fn command_stp0(&mut self, command: STP0Command, payload: &str) -> Result<String> {
        let message = format!("*{} {}", command.as_str(), payload);
        let command = format!("{} {}", message.as_str().len(), message);
        let encoded_command = UTF_16BE.encode(command.as_str(), EncoderTrap::Ignore).unwrap();

        self.send(&encoded_command)
    }

    fn wait_response(&mut self) -> Result<String> {
        let mut buf = [0u8; 2048];

        match self.socket.read(&mut buf) {
            Ok(_) => ScopeConnection::parse_response(buf),
            Err(e) => Err(e)
        }
    }

    fn parse_response(mut buf: [u8; 2048]) -> Result<String> {
        let terminator = b"\x20";

        match buf.iter().position(|&byte| byte == terminator[0]) {
            Some(pos) => {
                let (raw_size, raw) = buf.split_at_mut(pos);
                let size: u64 = UTF_16BE.decode(raw_size, DecoderTrap::Ignore).unwrap().parse::<u64>().unwrap() * 2;
                let (_, raw_message) = raw.split_at_mut(1);

                let mut message_buf: Vec<u8> = vec![];
                match raw_message.take(size).read_to_end(&mut message_buf) {
                    Ok(n) => {
                        let message = String::from(UTF_16BE.decode(&message_buf, DecoderTrap::Ignore).unwrap());
                        Ok(message)
                    },
                    Err(e) => Err(e)
                }
            },
            None => {
                let end = buf.iter().position(|&byte| byte == b'\n').unwrap();
                let (message, _) = buf.split_at_mut(end);

                match ASCII.decode(&message, DecoderTrap::Ignore) {
                    Ok(str) => {
                        Ok(String::from(str))
                    },
                    Err(_) => Err(Error::new(ErrorKind::InvalidInput, "Cannot decode input string"))
                }
            }
        }
    }

    fn send(&mut self, buf: &[u8]) -> Result<String> {
       match self.socket.write(buf) {
           Ok(_) => self.wait_response(),
           Err(e) => Err(e)
       }
    }

    fn parse_services_message(str: &String) -> Result<Vec<String>> {
        if str.contains("*services") {
            let services_str: String = String::from_str(str.split_whitespace().last().unwrap()).unwrap();

            Ok(services_str.split_terminator(",").map(String::from).collect())
        } else {
            Err(Error::new(ErrorKind::InvalidInput, "String must include a '*services' substring"))
        }
    }
}

fn handle_socket(socket: TcpStream) {
    let mut scope_connection = ScopeConnection { socket };

    match scope_connection.wait_response()
        .and_then(|response: String| ScopeConnection::parse_services_message(&response))
        .and_then(|services: Vec<String>| {
            println!("[services]:");
            match services.iter().position(|service| {
                println!("\t{}", service);
                service == "stp-1"
            }) {
                Some(_) => scope_connection.command_stp0(STP0Command::Enable, "stp-1"),
                None => Err(Error::new(ErrorKind::Other, "STP/1 is not supported"))
            }
        })
        .and_then(|response: String| {
            println!("[RECV] {}", response);

            scope_connection.wait_response()
        }) {
        Ok(response) => println!("[RECV] {}", response),
        Err(error) => println!("[ERR] {}", error)
    }
}

fn main() {
    let listener = TcpListener::bind("192.168.1.74:7001").unwrap();

    match listener.accept() {
        Ok((socket, addr)) => {
            println!("new client: {:?}", addr);
            handle_socket(socket);
        },
        Err(e) => println!("couldn't get client: {:?}", e),
    }
}