hightorrent 0.4.1

use bt_bencode::Value as BencodeValue;
use fluent_uri::pct_enc::{
    encoder::{Data, Query},
    EStr, EString,
};
use rustc_hex::ToHex;
#[cfg(feature = "sea_orm")]
use sea_orm::prelude::*;
use serde::{Deserialize, Serialize};
use sha1::{Digest, Sha1};

use std::collections::HashMap;
use std::path::PathBuf;

use crate::{
    InfoHash, InfoHashError, MagnetLink, MagnetLinkError, PieceLength, TorrentContent, TorrentID,
    Tracker,
};

/// Error occurred during parsing a [`TorrentFile`](crate::torrent_file::TorrentFile).
#[derive(Clone, Debug, PartialEq)]
pub enum TorrentFileError {
    NoNameFound,
    // TODO: bt_bencode::Error is not PartialEq so we store error as String
    InvalidBencode { reason: String },
    NotATorrent { reason: String },
    WrongVersion { version: u64 },
    InvalidHash { source: InfoHashError },
    InvalidContentPath { path: String },
    MissingPieceLength,
    BadPieceLength { piece_length: u32 },
}

impl std::fmt::Display for TorrentFileError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TorrentFileError::NoNameFound => write!(f, "No name found"),
            TorrentFileError::InvalidBencode { reason } => write!(f, "Invalid bencode: {reason}"),
            TorrentFileError::NotATorrent { reason } => write!(
                f,
                "Valid bencode, but does not seem to be a torrent ({reason})"
            ),
            TorrentFileError::WrongVersion { version } => write!(
                f,
                "Wrong torrent version: {version}, only v1 and v2 are supported)"
            ),
            TorrentFileError::InvalidHash { source } => write!(f, "Invalid hash: {source}"),
            TorrentFileError::InvalidContentPath { path } => {
                write!(f, "Invalid content file path in torrent: {path}")
            }
            TorrentFileError::MissingPieceLength => {
                write!(f, "No \'piece length\' field found in info dict")
            }
            TorrentFileError::BadPieceLength { piece_length } => {
                write!(f, "Torrent \'piece length\' is too big: {}", piece_length)
            }
        }
    }
}

impl From<InfoHashError> for TorrentFileError {
    fn from(e: InfoHashError) -> TorrentFileError {
        TorrentFileError::InvalidHash { source: e }
    }
}

impl From<bt_bencode::Error> for TorrentFileError {
    fn from(e: bt_bencode::Error) -> TorrentFileError {
        TorrentFileError::InvalidBencode {
            reason: e.to_string(),
        }
    }
}

impl std::error::Error for TorrentFileError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            TorrentFileError::InvalidHash { source } => Some(source),
            _ => None,
        }
    }
}

/// A torrent file.
///
/// The torrent file specification and related extensions are described on [Wikipedia](https://en.wikipedia.org/wiki/Torrent_file).
/// The TorrentFile can provide information about the torrent
/// [`name`](crate::torrent_file::TorrentFile::name) and
/// [`hash`](crate::torrent_file::TorrentFile::hash). Other fields could be supported, but are not
/// currently implemented by this library.
///
/// To save the torrent file to disk, use the `TorrentFile::to_vec` method:
///
/// ```ignore
/// std::fs::write("export.torrent", &torrent.to_vec()).unwrap();
/// ```
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct TorrentFile {
    pub hash: InfoHash,
    pub name: String,
    // Kept for further analysis
    pub decoded: DecodedTorrent,
}

/// A parsed bencode-decoded value, to ensure torrent-like structure.
///
/// In its present form, DecodedTorrent mostly cares about the info dict, but preserves other fields
/// as [`BencodeValue`](bt_bencode::BencodeValue) in an `extra` mapping so you can implement
/// your own extra parsing.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct DecodedTorrent {
    /// Main tracker
    #[serde(skip_serializing_if = "Option::is_none")]
    announce: Option<Tracker>,

    /// Many alternative trackers.
    /// TODO: what is this about the tiers?
    #[serde(
        rename = "announce-list",
        default,
        skip_serializing_if = "Vec::is_empty"
    )]
    announce_list: Vec<Vec<Tracker>>,

    info: DecodedInfo,

    // Rest of torrent dict
    #[serde(flatten)]
    #[serde(skip_serializing_if = "HashMap::is_empty")]
    extra: HashMap<String, BencodeValue>,
}

impl DecodedTorrent {
    pub fn files(&self) -> Result<Vec<TorrentContent>, TorrentFileError> {
        if let Some(info_files) = &self.info.files {
            // V1 torrent with multiple files
            let mut files: Vec<TorrentContent> = vec![];
            for file in info_files {
                // TODO: error
                let f: UnsafeV1FileContent = bt_bencode::from_value(file.clone()).unwrap();
                if let Some(parsed_file) = f.to_torrent_content()? {
                    files.push(parsed_file);
                }
            }

            // Sort files by alphabetical order
            files.sort();
            return Ok(files);
        }

        if let Some(_info_file_tree) = &self.info.file_tree {
            todo!("v2 torrent files");
        }

        // V1 torrent with single file
        Ok(vec![TorrentContent {
            path: PathBuf::from(&self.info.name),
            size: self.info.length.unwrap(),
        }])
    }
}

/// Raw file path described within a Bittorrent v1 torrent file.
///
/// It has not been sanitized, for example to prevent path traversal attacks. You should not be using this in your API;
/// use [TorrentContent] instead.
#[derive(Deserialize, Serialize, Debug, PartialEq, Clone)]
pub struct UnsafeV1FileContent {
    /// Raw path segments from the torrent, may contain directory escapes (like `..`)
    #[serde(rename = "path")]
    pub raw_paths: Vec<String>,
    /// File length in bytes
    pub length: u64,
    /// Extended file attributes as defined in [BEP-0047](https://www.bittorrent.org/beps/bep_0047.html)
    ///
    /// Can contain several characters:
    ///
    /// - p for padding files
    /// - l for symlinks
    /// - x for executables
    /// - h for hidden files
    #[serde(default)]
    pub attr: String,
}

impl UnsafeV1FileContent {
    /// Tries to parse [TorrentContent].
    ///
    /// Fails if the data is invalid (eg. path traversal), produces
    /// Ok(None) when the file is a padding file.
    pub fn to_torrent_content(&self) -> Result<Option<TorrentContent>, TorrentFileError> {
        if self.attr.contains('p') {
            return Ok(None);
        }

        // Parse the raw path parts omitting weird directory shenanigans
        let mut path = PathBuf::new();
        for p in &self.raw_paths {
            if p.contains('/') {
                return Err(TorrentFileError::InvalidContentPath {
                    path: p.to_string(),
                });
            }

            if p == ".." {
                return Err(TorrentFileError::InvalidContentPath {
                    path: p.to_string(),
                });
            }

            if p == "." {
                continue;
            }

            path.push(p);
        }

        Ok(Some(TorrentContent {
            path,
            size: self.length,
        }))
    }
}

/// An info dict contained in a [`DecodedTorrent`](crate::torrent_file::DecodedTorrent).
///
/// Only cares about torrent version, name, and files, but other fields are preseved in an `extra`
/// mapping so you can implement your own extra parsing.
// bt_bencode does not support serializing None options and empty HashMaps, so we skip
// serialization in those cases.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct DecodedInfo {
    #[serde(rename = "meta version")]
    #[serde(skip_serializing_if = "Option::is_none")]
    version: Option<u64>,

    name: String,

    /// Torrent `piece length` as used in v1/v2 torrents
    ///
    /// Maximum size is 536854528 like in libtorrent.
    #[serde(rename = "piece length")]
    piece_length: PieceLength,

    // Torrent v1/hybrid (only for single-file torrents)
    #[serde(skip_serializing_if = "Option::is_none")]
    length: Option<u64>,

    // Torrent v1 (only for multi-files torrents)
    #[serde(skip_serializing_if = "Option::is_none")]
    files: Option<Vec<BencodeValue>>,

    // Torrent v2 (for both single and multi-files torrents)
    #[serde(rename = "file tree")]
    #[serde(skip_serializing_if = "Option::is_none")]
    file_tree: Option<BencodeValue>,

    // Rest of info dict that we keep for hashing
    #[serde(flatten)]
    #[serde(skip_serializing_if = "HashMap::is_empty")]
    extra: HashMap<String, BencodeValue>,
}

impl TorrentFile {
    /// Serialize to a .torrent file byte slice
    pub fn to_vec(&self) -> Vec<u8> {
        // This should not fail
        bt_bencode::to_vec(&self.decoded).unwrap()
    }

    /// Deserialize (parse) from a .torrent file byte slice
    pub fn from_slice(s: &[u8]) -> Result<TorrentFile, TorrentFileError> {
        let torrent: DecodedTorrent = bt_bencode::from_slice(s).map_err(|e| {
            // We store a stringy representation of the error because bt_encode::Error
            // is not PartialEq
            TorrentFileError::NotATorrent {
                reason: e.to_string(),
            }
        })?;

        // We just deserialized successfully so this is a safe unwrap
        // Unless we added an Option/HashMap and forgot to skip serialization when empty
        let info_bytes = bt_bencode::to_vec(&torrent.info).unwrap();

        let infohash = match torrent.info.version {
            // Most v1 torrents don't declare a torrent version at all
            Some(1) | None => {
                // Bittorrent v1 does not necessarily have a files dict... single-file torrents
                // just use the torrent name field for that
                let digest = Sha1::digest(&info_bytes).to_vec().to_hex::<String>();
                InfoHash::new(&digest)?
            }
            Some(2) => {
                // Bittorrent v2 has mandatory file_tree dict
                // see http://bittorrent.org/beps/bep_0052.html
                if torrent.info.file_tree.is_some() {
                    let digest = sha256::digest(info_bytes.as_slice());
                    let hash = InfoHash::new(&digest)?;
                    // Check if we have hybrid torrent...
                    // If it's single-file it will have length field
                    // If it's multi-file it will have files field
                    if torrent.info.length.is_some() || torrent.info.files.is_some() {
                        let digest = Sha1::digest(&info_bytes).to_vec().to_hex::<String>();
                        hash.hybrid(&InfoHash::new(&digest)?)?
                    } else {
                        hash
                    }
                } else {
                    return Err(TorrentFileError::NotATorrent {
                        reason: "Torrentv2 without 'file_tree' field".to_string(),
                    });
                }
            }
            _ => {
                // Version is not null and is not 1-2
                return Err(TorrentFileError::WrongVersion {
                    version: torrent.info.version.unwrap(),
                });
            }
        };

        Ok(TorrentFile {
            name: torrent.info.name.clone(),
            hash: infohash,
            decoded: torrent,
        })
    }

    pub fn hash(&self) -> &str {
        self.hash.as_str()
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn id(&self) -> TorrentID {
        TorrentID::from_infohash(&self.hash)
    }

    /// List the trackers in the torrent, ensuring
    /// they're sorted and deduplicated.
    pub fn trackers(&self) -> Vec<Tracker> {
        let mut trackers = vec![];
        if let Some(tracker) = &self.decoded.announce {
            trackers.push(tracker.clone());
        }

        for tier in &self.decoded.announce_list {
            for tracker in tier {
                trackers.push(tracker.clone());
            }
        }

        trackers.sort_unstable();
        trackers.dedup();

        trackers
    }

    /// Lossy transformation into a [`MagnetLink`].
    ///
    /// This is a lossy operation because:
    ///
    /// - magnet links do not contain detailed pieces information,
    ///   so there is no reverse operation without network resolution
    /// - not all metadata about the torrent may be added to the magnet link
    ///
    /// What's added to the magnet link:
    ///
    /// - name
    /// - torrentID
    /// - trackers
    ///
    /// For the moment, this is a fallible operation because we make sure
    /// the produced MagnetLink can be parsed again. This operation may not
    /// produce errors in a future release.
    pub fn magnet_link(&self) -> Result<MagnetLink, MagnetLinkError> {
        // Not sure how to build a URI with `magnet:` and not `magnet://`
        // without the `Uri::builder`.
        let mut uri = match &self.hash {
            InfoHash::V1(h) => format!("magnet:?xt=urn:btih:{h}"),
            InfoHash::V2(h) => format!("magnet:?xt=urn:btmh:1220{h}"),
            InfoHash::Hybrid((h1, h2)) => format!("magnet:?xt=urn:btih:{h1}&xt=urn:btmh:1220{h2}"),
        };

        let mut buf = EString::<Query>::new();

        if !self.name.is_empty() {
            buf.push_estr(EStr::new_or_panic("&dn="));
            buf.encode_str::<Data>(&self.name);
        }

        // We use the helper to avoid duplicates or unsorted entries
        for tracker in self.trackers() {
            buf.push_estr(EStr::new_or_panic("&tr="));
            buf.encode_str::<Data>(tracker.url());
        }

        uri.push_str(buf.as_str());

        MagnetLink::new(&uri)
    }
}

#[cfg(feature = "sea_orm")]
impl From<TorrentFile> for sea_orm::sea_query::Value {
    fn from(t: TorrentFile) -> Self {
        Value::Bytes(Some(t.to_vec()))
    }
}

#[cfg(feature = "sea_orm")]
impl sea_orm::TryGetable for TorrentFile {
    fn try_get_by<I: sea_orm::ColIdx>(
        res: &sea_orm::QueryResult,
        index: I,
    ) -> Result<Self, sea_orm::error::TryGetError> {
        let val = <Vec<u8> as sea_orm::TryGetable>::try_get_by(res, index)?;
        TorrentFile::from_slice(&val).map_err(|e| {
            sea_orm::error::TryGetError::DbErr(sea_orm::DbErr::TryIntoErr {
                from: "Bytes",
                into: "TorrentFile",
                source: std::sync::Arc::new(e),
            })
        })
    }
}

#[cfg(feature = "sea_orm")]
impl sea_orm::sea_query::ValueType for TorrentFile {
    fn try_from(v: sea_orm::Value) -> Result<Self, sea_orm::sea_query::ValueTypeErr> {
        match v {
            sea_orm::Value::Bytes(Some(s)) => {
                TorrentFile::from_slice(&s).map_err(|_e| sea_orm::sea_query::ValueTypeErr)
            }
            _ => Err(sea_orm::sea_query::ValueTypeErr),
        }
    }

    fn type_name() -> String {
        "TorrentFile".to_string()
    }

    fn array_type() -> sea_orm::sea_query::ArrayType {
        sea_orm::sea_query::ArrayType::Bytes
    }

    fn column_type() -> sea_orm::sea_query::ColumnType {
        sea_orm::sea_query::ColumnType::VarBinary(StringLen::None)
    }
}

#[cfg(feature = "sea_orm")]
impl sea_orm::sea_query::Nullable for TorrentFile {
    fn null() -> sea_orm::sea_query::Value {
        sea_orm::sea_query::Value::Bytes(None)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn can_read_torrent_v1_multifile() {
        let slice = std::fs::read("tests/bittorrent-v1-emma-goldman.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        println!("{:?}", res);
        assert!(res.is_ok());
        let torrent = res.unwrap();
        assert_eq!(
            &torrent.name,
            "Goldman, Emma - Essential Works of Anarchism"
        );
        assert_eq!(
            torrent.hash,
            InfoHash::V1("c811b41641a09d192b8ed81b14064fff55d85ce3".to_string())
        );
        assert_eq!(torrent.decoded.files().unwrap().len(), 94);
    }

    #[test]
    #[cfg(not(feature = "unknown_tracker_scheme"))]
    fn fail_no_torrent_scheme() {
        let slice = std::fs::read("tests/libtorrent/good/sample.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        println!("{:?}", res);
        assert!(res.is_err());
        assert_eq!(
            res.unwrap_err(),
            TorrentFileError::NotATorrent {
                reason: "Invalid scheme: tracker.publicbt.com".to_string()
            },
        );
    }

    #[test]
    fn can_read_torrent_v1_wrongpath() {
        let slice = std::fs::read("tests/libtorrent/good/parent_path.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        println!("{:?}", res);
        assert!(res.is_ok());
        let torrent = res.unwrap();
        assert_eq!(&torrent.name, "temp");
        assert_eq!(
            torrent.hash,
            InfoHash::V1("9e1111f1ee4966f7d06d398f1d58e00ad150657a".to_string())
        );
        assert_eq!(
            torrent.decoded.files().unwrap_err(),
            TorrentFileError::InvalidContentPath {
                path: "..".to_string()
            },
        );
    }

    #[test]
    fn can_read_torrent_v1_singlepath() {
        let slice = std::fs::read("tests/libtorrent/good/base.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        println!("{:?}", res);
        assert!(res.is_ok());
        let torrent = res.unwrap();
        assert_eq!(&torrent.name, "temp");
        assert_eq!(
            torrent.hash,
            InfoHash::V1("c0fda1edafdbdbb96443424e0b3899af7159d10e".to_string())
        );
        assert_eq!(
            torrent.decoded.files().unwrap(),
            vec!(TorrentContent {
                path: PathBuf::from("temp"),
                size: 425,
            }),
        );
    }

    #[test]
    fn can_read_torrent_v2() {
        let slice = std::fs::read("tests/bittorrent-v2-test.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        assert!(res.is_ok());
        let torrent = res.unwrap();
        assert_eq!(&torrent.name, "bittorrent-v2-test");
        assert_eq!(
            torrent.hash,
            InfoHash::V2(
                "caf1e1c30e81cb361b9ee167c4aa64228a7fa4fa9f6105232b28ad099f3a302e".to_string()
            )
        );
    }

    #[test]
    fn can_read_torrent_hybrid() {
        let slice = std::fs::read("tests/bittorrent-v2-hybrid-test.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        assert!(res.is_ok());
        let torrent = res.unwrap();
        assert_eq!(&torrent.name, "bittorrent-v1-v2-hybrid-test");
        assert_eq!(
            torrent.hash,
            InfoHash::Hybrid((
                "631a31dd0a46257d5078c0dee4e66e26f73e42ac".to_string(),
                "d8dd32ac93357c368556af3ac1d95c9d76bd0dff6fa9833ecdac3d53134efabb".to_string()
            ))
        );
    }

    #[test]
    fn v1_piece_len() {
        let slice = std::fs::read("tests/libtorrent/bad/negative_piece_len.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        assert!(res.is_err());
    }

    #[test]
    fn v2_piece_len() {
        let slice = std::fs::read("tests/libtorrent/bad/v2_piece_size.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        assert!(res.is_err());
    }

    #[test]
    fn test_torrent_to_magnet_v2() {
        let slice = std::fs::read("tests/bittorrent-v2-test.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        let torrent = res.unwrap();

        let expected = std::fs::read_to_string("tests/bittorrent-v2-test.magnet").unwrap();
        let magnet = torrent.magnet_link().unwrap();
        assert_eq!(expected, magnet.to_string(),);
    }

    #[test]
    fn test_torrent_to_magnet_hybrid() {
        let slice = std::fs::read("tests/bittorrent-v2-hybrid-test.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        let torrent = res.unwrap();

        let expected = std::fs::read_to_string("tests/bittorrent-v2-hybrid-test.magnet").unwrap();
        let magnet = torrent.magnet_link().unwrap();
        assert_eq!(expected, magnet.to_string(),);
    }

    #[test]
    fn test_torrent_to_magnet_v1() {
        let slice = std::fs::read("tests/bittorrent-v1-emma-goldman.torrent").unwrap();
        let res = TorrentFile::from_slice(&slice);
        let torrent = res.unwrap();

        let expected = MagnetLink::new(
            &std::fs::read_to_string("tests/bittorrent-v1-emma-goldman.magnet").unwrap(),
        )
        .unwrap();
        let magnet = torrent.magnet_link().unwrap();

        assert_eq!(expected.name(), magnet.name(),);

        assert_eq!(expected.hash(), magnet.hash(),);

        // Check for duplicates. This is useful because the `announce` in a torrent
        // file may also be in the `announce_list` so we wanna make sure we don't have it twice.
        for tracker in magnet.trackers() {
            if magnet.trackers().iter().filter(|x| x == &tracker).count() > 1 {
                panic!("Duplicate tracker: {tracker:?}");
            }
        }

        // Check the trackers are actually equal
        assert_eq!(magnet.trackers(), expected.trackers());

        // Check for complete equality in this specific case (no information loss on the test)
        assert_eq!(expected, magnet);
    }
}