path: root/src/client/encrypted.rs

use std::str;

use quick_xml::{
    events::{BytesDecl, Event},
    name::QName,
    Reader, Writer,
};
use rsasl::prelude::{Mechname, SASLClient};
use tokio::io::{BufReader, ReadHalf, WriteHalf};
use tokio::net::TcpStream;
use tokio_native_tls::TlsStream;

use crate::stanza::{
    bind::Bind,
    iq::IQ,
    sasl::{Challenge, Success},
    Element,
};
use crate::stanza::{
    sasl::{Auth, Response},
    stream::{Stream, StreamFeature},
};
use crate::Jabber;
use crate::JabberError;
use crate::Result;

pub struct JabberClient<'j> {
    pub reader: Reader<BufReader<ReadHalf<TlsStream<TcpStream>>>>,
    pub writer: Writer<WriteHalf<TlsStream<TcpStream>>>,
    jabber: &'j mut Jabber<'j>,
}

impl<'j> JabberClient<'j> {
    pub fn new(
        reader: Reader<BufReader<ReadHalf<TlsStream<TcpStream>>>>,
        writer: Writer<WriteHalf<TlsStream<TcpStream>>>,
        jabber: &'j mut Jabber<'j>,
    ) -> Self {
        Self {
            reader,
            writer,
            jabber,
        }
    }

    pub async fn start_stream(&mut self) -> Result<()> {
        // client to server
        let declaration = BytesDecl::new("1.0", None, None);
        let server = &self.jabber.server.to_owned().try_into()?;
        let stream_element =
            Stream::new_client(&self.jabber.jid, server, None, Some("en".to_string()));
        self.writer
            .write_event_async(Event::Decl(declaration))
            .await;
        let stream_element: Element<'_> = stream_element.into();
        stream_element.write_start(&mut self.writer).await?;
        // server to client
        let mut buf = Vec::new();
        self.reader.read_event_into_async(&mut buf).await?;
        let _stream_response = Element::read_start(&mut self.reader).await?;
        Ok(())
    }

    pub async fn get_features(&mut self) -> Result<Vec<StreamFeature>> {
        Element::read(&mut self.reader).await?.try_into()
    }

    pub async fn negotiate(&mut self) -> Result<()> {
        loop {
            println!("negotiate loop");
            let features = self.get_features().await?;
            println!("features: {:?}", features);

            match &features[0] {
                StreamFeature::Sasl(sasl) => {
                    println!("sasl?");
                    self.sasl(&sasl).await?;
                }
                StreamFeature::Bind => {
                    self.bind().await?;
                    return Ok(());
                }
                x => println!("{:?}", x),
            }
        }
    }

    pub async fn watch(&mut self) -> Result<()> {
        loop {
            let element = Element::read(&mut self.reader).await?;
            println!("{:#?}", element);
        }
    }

    pub async fn sasl(&mut self, mechanisms: &Vec<String>) -> Result<()> {
        println!("{:?}", mechanisms);
        let sasl = SASLClient::new(self.jabber.auth.clone());
        let mut offered_mechs: Vec<&Mechname> = Vec::new();
        for mechanism in mechanisms {
            offered_mechs.push(Mechname::parse(mechanism.as_bytes())?)
        }
        println!("{:?}", offered_mechs);
        let mut session = sasl.start_suggested(&offered_mechs)?;
        let selected_mechanism = session.get_mechname().as_str().to_owned();
        println!("selected mech: {:?}", selected_mechanism);
        let mut data: Option<Vec<u8>> = None;
        if !session.are_we_first() {
            // if not first mention the mechanism then get challenge data
            // mention mechanism
            let auth = Auth {
                mechanism: selected_mechanism.as_str(),
                sasl_data: "=",
            };
            Into::<Element>::into(auth).write(&mut self.writer).await?;
            // get challenge data
            let challenge = &Element::read(&mut self.reader).await?;
            let challenge: Challenge = challenge.try_into()?;
            println!("challenge: {:?}", challenge);
            data = Some(challenge.sasl_data.to_owned());
            println!("we didn't go first");
        } else {
            // if first, mention mechanism and send data
            let mut sasl_data = Vec::new();
            session.step64(None, &mut sasl_data).unwrap();
            let auth = Auth {
                mechanism: selected_mechanism.as_str(),
                sasl_data: str::from_utf8(&sasl_data)?,
            };
            println!("{:?}", auth);
            Into::<Element>::into(auth).write(&mut self.writer).await?;

            let server_response = Element::read(&mut self.reader).await?;
            println!("server_response: {:#?}", server_response);
            match TryInto::<Challenge>::try_into(&server_response) {
                Ok(challenge) => data = Some(challenge.sasl_data.to_owned()),
                Err(_) => {
                    let success = TryInto::<Success>::try_into(&server_response)?;
                    if let Some(sasl_data) = success.sasl_data {
                        data = Some(sasl_data.to_owned())
                    }
                }
            }
            println!("we went first");
        }

        // stepping the authentication exchange to completion
        if data != None {
            println!("data: {:?}", data);
            let mut sasl_data = Vec::new();
            while {
                // decide if need to send more data over
                let state = session
                    .step64(data.as_deref(), &mut sasl_data)
                    .expect("step errored!");
                state.is_running()
            } {
                // While we aren't finished, receive more data from the other party
                let response = Response {
                    sasl_data: str::from_utf8(&sasl_data)?,
                };
                println!("response: {:?}", response);
                Into::<Element>::into(response)
                    .write(&mut self.writer)
                    .await?;

                let server_response = Element::read(&mut self.reader).await?;
                println!("server_response: {:?}", server_response);
                match TryInto::<Challenge>::try_into(&server_response) {
                    Ok(challenge) => data = Some(challenge.sasl_data.to_owned()),
                    Err(_) => {
                        let success = TryInto::<Success>::try_into(&server_response)?;
                        if let Some(sasl_data) = success.sasl_data {
                            data = Some(sasl_data.to_owned())
                        }
                    }
                }
            }
        }
        self.start_stream().await?;
        Ok(())
    }

    pub async fn bind(&mut self) -> Result<()> {
        match &self.jabber.jid.resourcepart {
            Some(resource) => {
                println!("setting resource");
                let bind = Bind {
                    resource: Some(resource.clone()),
                    jid: None,
                };
                let result: Bind = IQ::set(self, None, None, bind).await?.try_into()?;
                if let Some(jid) = result.jid {
                    println!("{}", jid);
                    self.jabber.jid = jid;
                    return Ok(());
                }
            }
            None => {
                println!("not setting resource");
                let bind = Bind {
                    resource: None,
                    jid: None,
                };
                let result: Bind = IQ::set(self, None, None, bind).await?.try_into()?;
                if let Some(jid) = result.jid {
                    println!("{}", jid);
                    self.jabber.jid = jid;
                    return Ok(());
                }
            }
        }
        Err(JabberError::BindError)
    }
}