extern crate rand;
extern crate tokio;
extern crate openssl;
extern crate byteorder;

use std::net::{IpAddr, Ipv4Addr, SocketAddr};

use rand::{thread_rng, Rng};

use tokio::io;
use tokio::prelude::*;
use tokio::net::{TcpListener, TcpStream};

use openssl::sha::sha256;
use openssl::symm::{encrypt, decrypt, Cipher};

use byteorder::{LittleEndian, BigEndian, ReadBytesExt, WriteBytesExt};

// const fn has not been stabilished yet
// static + custom From is the closest we can get
#[derive(Debug, Clone, Copy)]
struct IpAddress(u8, u8, u8, u8);

impl From<IpAddress> for IpAddr {
    fn from(addr: IpAddress) -> Self {
        IpAddr::V4(Ipv4Addr::new(addr.0, addr.1, addr.2, addr.3))
    }
}

const DC_IPS: [IpAddress; 5] = [
    IpAddress(149, 154, 175,  50),
    IpAddress(149, 154, 167,  51),
    IpAddress(149, 154, 175, 100),
    IpAddress(149, 154, 167,  91),
    IpAddress(149, 154, 171,   5),
];

struct MTProtoPacket {
    encrypt_key: Vec<u8>,
    encrypt_iv:  Vec<u8>,
    decrypt_key: Vec<u8>,
    decrypt_iv:  Vec<u8>,

    encrypt_counter: Option<Vec<u8>>,
    decrypt_counter: Option<Vec<u8>>,

    //encrypt_count_buf: Vec<u8>
    //encrypt_num: u32,

    //decrypt_count_buf: Vec<u8>,
    //decrypt_buf: u32,
}

impl MTProtoPacket {
    fn new(buffer: &Vec<u8>, rev: &Vec<u8>,
           key: &Vec<u8>, key_rev: &Vec<u8>) -> MTProtoPacket {
        // This is verified
        MTProtoPacket {
            encrypt_key: key_rev.clone(),
            encrypt_iv: rev[32..48].to_vec(),
            decrypt_key: key.clone(),
            decrypt_iv: buffer[40..56].to_vec(),
            encrypt_counter: None,
            decrypt_counter: None
        }
    }

    /// Initializes a new obfuscated2 buffer
    fn new_obf2_buffer() -> (MTProtoPacket, [u8; 64]) {
        let mut buffer = [0u8; 64];
        let mut rng = thread_rng();
        loop {
            rng.fill_bytes(&mut buffer);

            let val = (&buffer[0..4]).read_u32::<LittleEndian>().unwrap();
            let val2 = (&buffer[4..8]).read_u32::<LittleEndian>().unwrap();
            if buffer[0] != 0xef
                && val != 0x44414548
                && val != 0x54534f50
                && val != 0x20544547
                && val != 0x4954504f
                && val != 0xeeeeeeee
                && val2 != 0x00000000
            {
                buffer[56] = 0xef;
                buffer[57] = 0xef;
                buffer[58] = 0xef;
                buffer[59] = 0xef;
                break;
            }
        }

        let encrypt_key;
        let encrypt_iv;
        let decrypt_key;
        let decrypt_iv;
        {
            let key_iv_encrypt: &[u8] = &buffer[8..56];
            encrypt_key = key_iv_encrypt[0..32].to_vec();
            encrypt_iv = key_iv_encrypt[32..48].to_vec();

            let key_iv_encrypt: Vec<u8> = key_iv_encrypt.iter().rev().cloned().collect();
            decrypt_key = key_iv_encrypt[0..32].to_vec();
            decrypt_iv = key_iv_encrypt[32..48].to_vec();
        }

        let mut packet = MTProtoPacket {
            encrypt_key, encrypt_iv,
            decrypt_key, decrypt_iv,
            encrypt_counter: None,
            decrypt_counter: None
        };
        let encrypted_buffer = packet.encrypt_obf2(&buffer);
        for i in 56..encrypted_buffer.len() {
            buffer[i] = encrypted_buffer[i];
        }

        (packet, buffer)
    }

    fn encrypt_obf2(&mut self, data: &[u8]) -> Vec<u8> {
        //if self.encrypt_counter.is_none() {
        //    self.encrypt_counter = Some(self.encrypt_iv.cloned());
        //}

        let mut out = vec![0u8; data.len()];
        let cipher = Cipher::aes_256_ctr();
        encrypt(cipher, &self.encrypt_key, Some(&self.encrypt_iv), &data).unwrap()
    }

    fn decrypt_obf2(&mut self, data: &[u8]) -> Vec<u8> {
        //if self.decrypt_counter.is_none() {
        //    self.decrypt_counter = Some(self.decrypt_iv.cloned());
        //}

        let mut out = vec![0u8; data.len()];
        let cipher = Cipher::aes_256_ctr();
        decrypt(cipher, &self.decrypt_key, Some(&self.decrypt_iv), &data).unwrap()
    }

    fn create_packet_obf2(&mut self, payload: &[u8]) -> Vec<u8> {
        // TODO lock?
        let payload = self.create_packet_abridged(&payload);
        self.encrypt_obf2(&payload);
    }

    fn create_packet_abridged(&self, payload: &[u8]) -> Vec<u8> {
        let length = payload.len() / 4;
        let mut bytes = Vec::new();
        if length < 0x7f {
            bytes.push(length as u8);
        } else {
            bytes.push(0x7f);
            bytes.write_u32::<LittleEndian>(length as u32).unwrap();
            bytes.remove(1);
        }
        bytes.extend(payload);
        bytes
    }

    fn send_init_buffer_obf2() {
        // TODO wtf
    }

    // this is GetPacketBytesAsync
    fn read(&mut self, socket: TcpStream) {
        let connection = io::read_exact(socket, [0u8; 1])
            .and_then(|(socket, buf)| {
                let len = self.decrypt_obf2(&buf)[0];
                if len < 0x7f {
                    Ok((socket, len * 4))
                } else {
                    io::read_exact(socket, [0u8; 3])
                        .and_then(|socket, buf|) {
                            let mut len_bytes = self.decrypt_obf2(&buf);
                            len_bytes.insert(0, 0);
                            Ok((socket, len_bytes.read_i32::<LittleEndian>().unwrap()))
                        }
                }
                Ok(io::read_exact(socket, [0u8; 3]))
            });

    }
}

struct MTProtoProxyServer {
    secret: u128,
    port: u16
}

fn to_bytes(n: u128) -> Vec<u8> {
    let mut vec = vec![];
    vec.write_u128::<BigEndian>(n).unwrap();
    vec
}

impl MTProtoProxyServer {
    fn new(secret: u128, port: u16) -> MTProtoProxyServer {
        MTProtoProxyServer { secret, port }
    }

    fn run(&self) {
        let addr = format!("192.168.1.104:{}", self.port).parse().unwrap();
        let listener = TcpListener::bind(&addr).expect("bind fail");
        let secret = self.secret;

        let server = listener.incoming().for_each(move |socket| {
            println!("got socket at {:?}", socket.peer_addr().unwrap());

            let mut random_buffer = vec![0u8; 64];
            let connection = io::read_exact(socket, random_buffer)
                .and_then(move |(socket, mut buf)| {
                    let rev: Vec<u8> = buf[8..56].iter().rev().cloned().collect();
                    let key: Vec<u8> = buf[8..40].to_vec();
                    let key_rev: Vec<u8> = rev[0..32].to_vec();
                    let bin_secret: Vec<u8> = to_bytes(secret);

                    let mut key = key.to_vec();
                    key.extend(&bin_secret);
                    let key = sha256(&key);

                    let mut key_rev = key_rev.to_vec();
                    key_rev.extend(&bin_secret);
                    let key_rev = sha256(&key_rev);

                    let mut mtproto_packet_sv = MTProtoPacket::new(
                        &buf, &rev, &key.to_vec(), &key_rev.to_vec()
                    );
                    let decrypt_buf = &mtproto_packet_sv.decrypt_obf2(&buf)[56..64];
                    for i in 56..64 {
                        buf[i] = decrypt_buf[i - 56];
                    }
                    for byte in &buf[56..60] {
                        if *byte != 0xefu8 {
                            eprintln!("error in buffer");
                            // TODO return Err
                        }
                    }

                    let dc_id = (&buf[60..62]).read_i16::<LittleEndian>().unwrap().abs();
                    let ip = DC_IPS[(dc_id - 1) as usize];
                    println!("connecting to dc {}: {:?}", dc_id, ip);

                    // ALL ABOVE VERIFIED UP TO HERE
                    
                    let socket_cl = TcpStream::connect(&SocketAddr::new(From::from(ip), 443));
                    let (mtproto_packet_cl, buffer) =
                        MTProtoPacket::new_obf2_buffer();

                    // TODO Actually don't but yeah testing!
                    return io::write_all(socket, buf);
                })
                .then(|_| {
                    Ok(())
                });

            tokio::spawn(connection);
            Ok(())
        })
        .map_err(|err| {
            eprintln!("acept error: {:?}", err);
        });

        println!("server running at {:?}", addr);
        tokio::run(server);
    }
}

fn main() {
    let server = MTProtoProxyServer::new(1837, 1837);
    println!("ok running...");
    server.run()
}