lena/crates/transmission-rs/src/client.rs

//! Client for download management.

use std::collections::BTreeMap;
use std::convert::TryInto;
use std::time::Duration;

use serde::Serialize;

use tokio::time::sleep;

use crate::Error;
mod payload;
use payload::Request;
use payload::Response;
mod list;
use list::FieldList;
use list::TorrentList;
pub mod torrent;
use torrent::AddTorrent;
use torrent::AddTorrentResponse;
use torrent::GetTorrentFilesRequest;
use torrent::GetTorrentFilesResponse;
use torrent::Torrent;
use torrent::TorrentFile;

/// Interface into the major functions of Transmission
/// including adding, and removing torrents.
///
/// The `Client` does not keep track of the created torrents itself.
/// ```
#[derive(Clone)]
pub struct Client {
    host: String,
    port: u16,
    tls: bool,
    auth: Option<(String, String)>,

    base_url: String,
    session_id: Option<String>,

    http: reqwest::Client,
}

impl Client {
    pub fn new<T, U>(
        host: impl ToString,
        port: u16,
        tls: bool,
        auth: Option<(T, U)>,
    ) -> Result<Self, Error>
    where
        T: Into<String>,
        U: Into<String>,
    {
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert(
            reqwest::header::USER_AGENT,
            "transmission-rs/0.1".try_into()?,
        );
        let auth = match auth {
            Some((u, p)) => Some((u.into(), p.into())),
            None => None,
        };
        let this = Self {
            host: host.to_string(),
            port,
            tls,
            auth,
            base_url: {
                let protocol = match tls {
                    true => "https",
                    false => "http",
                };
                format!("{}://{}:{}", protocol, host.to_string(), port)
            },
            session_id: None,
            http: reqwest::Client::builder()
                .gzip(true)
                .default_headers(headers)
                .build()?,
        };
        Ok(this)
    }

    fn build_url(&self, path: impl AsRef<str>) -> String {
        format!("{}{}", self.base_url, path.as_ref()).to_string()
    }

    pub async fn authenticate(&mut self) -> Result<(), Error> {
        self.session_id = None;
        let response = self
            .post("/transmission/rpc/")
            .header("Content-Type", "application/x-www-form-urlencoded")
            .send()
            .await?;
        self.session_id = match response.headers().get("X-Transmission-Session-Id") {
            Some(v) => Some(v.to_str()?.to_string()),
            None => None, // TODO:  Return an error
        };
        Ok(())
    }

    pub fn post(&self, path: impl AsRef<str>) -> reqwest::RequestBuilder {
        let url = self.build_url(path);
        //println!("{}", url);
        let mut request = self.http.post(&url);
        request = match &self.auth {
            Some(auth) => request.basic_auth(&auth.0, Some(&auth.1)),
            None => request,
        };
        request = match &self.session_id {
            Some(token) => request.header("X-Transmission-Session-Id", token),
            None => request,
        };
        request
    }

    pub async fn rpc(
        &mut self,
        method: impl ToString,
        tag: u8,
        body: impl Serialize,
    ) -> Result<reqwest::Response, Error> {
        let request = Request::new(method, tag, body);
        let mut error = None;
        for i in 0..5 {
            sleep(Duration::from_secs(i)).await;
            let result = self
                .post("/transmission/rpc/")
                // Content-Type doesn't actually appear to be necessary, and is
                //    technically a lie, since there's no key/value being passed,
                //    just some JSON, but the official client sends this, so we
                //    will too.
                .header("Content-Type", "application/x-www-form-urlencoded")
                .body(serde_json::to_string(&request)?)
                .send()
                .await?
                .error_for_status();
            match result {
                Ok(r) => return Ok(r),
                Err(e) => match e.status() {
                    Some(reqwest::StatusCode::CONFLICT) => {
                        self.authenticate().await?;
                        error = Some(e);
                        continue;
                    }
                    _ => return Err(e.into()),
                },
            };
        }
        match error {
            Some(e) => Err(e.into()),
            // Should be unreachable
            None => Err(Error::HTTPUnknown),
        }
    }

    pub async fn list(&mut self) -> Result<Vec<Torrent>, Error> {
        let field_list = FieldList::from_vec(vec![
            "error",
            "errorString",
            "eta",
            "id",
            "isFinished",
            "leftUntilDone",
            "name",
            "peersGettingFromUs",
            "peersSendingToUs",
            "rateDownload",
            "rateUpload",
            "sizeWhenDone",
            "status",
            "uploadRatio",
        ]);
        let response: Response<TorrentList> = self
            .rpc("torrent-get", 4, field_list)
            .await?
            .error_for_status()?
            .json()
            .await?;
        Ok(response.arguments.torrents)
    }

    // TODO:  Borrow key from value
    pub async fn list_by_name(&mut self) -> Result<BTreeMap<String, Torrent>, Error> {
        Ok(self
            .list()
            .await?
            .into_iter()
            .map(|v| (v.name.clone(), v))
            .collect::<BTreeMap<_, _>>())
    }

    pub async fn add_torrent(&mut self, data: AddTorrent) -> Result<AddTorrentResponse, Error> {
        let response: Response<AddTorrentResponse> = self
            .rpc("torrent-add", 8, data)
            .await?
            .error_for_status()?
            .json()
            .await?;
        Ok(response.arguments)
    }

    pub async fn get_files_by_id(&mut self, id: u16) -> Result<Vec<TorrentFile>, Error> {
        let response: Response<GetTorrentFilesResponse> = self
            .rpc("torrent-get", 3, GetTorrentFilesRequest::new(id))
            .await?
            .error_for_status()?
            .json()
            .await?;
        if (response.arguments.torrents.len() == 0) {
            // TODO:  Maybe make the Ok result an Option?  Or maybe return a 404 error?
            return Ok(vec![]);
        } else if (response.arguments.torrents.len() > 1) {
            return Err(Error::WeirdGetFilesResponse);
        }
        // TODO:  Figure out how to move files out of this structure, since we're discarding the rest anyway, instead of cloning
        Ok(response.arguments.torrents[0].files.clone())
    }
}