torment/torment-core/src/metainfo.rs

use crate::infohash::v1::U160;
use crate::infohash::InfoHashCapable;
use bytes::Bytes;
use num_traits::cast::ToPrimitive;
use ring::digest::{digest, SHA1_FOR_LEGACY_USE_ONLY};
use serde::export::fmt::Debug;
use serde::export::Formatter;
use std::collections::HashSet;
use std::fmt::Display;
use std::string::FromUtf8Error;
use torment_bencode::{BencodeError, BencodeValue};
use url::Url;

#[derive(Debug)]
pub enum MetaInfoParsingError {
    DecodingFailure(BencodeError),
    Utf8Error(FromUtf8Error),
    WrongType,
    MissingEntry,
}

impl From<BencodeError> for MetaInfoParsingError {
    fn from(err: BencodeError) -> Self {
        MetaInfoParsingError::DecodingFailure(err)
    }
}

impl From<std::string::FromUtf8Error> for MetaInfoParsingError {
    fn from(err: FromUtf8Error) -> Self {
        MetaInfoParsingError::Utf8Error(err)
    }
}

impl Display for MetaInfoParsingError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

#[derive(Clone)]
pub struct Torrent {
    name: String,
    announce_list: Vec<HashSet<Url>>,
    info: MetaInfo,
    bencoded: BencodeValue,
}

impl Debug for Torrent {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Torrent")
            .field("name", &self.name)
            .field("announce_list", &self.announce_list)
            .field("info", &self.info)
            .finish()
    }
}

impl Torrent {
    pub fn from_bytes<T: AsRef<[u8]>>(
        bytes: T,
        name: Option<String>,
    ) -> Result<Self, MetaInfoParsingError> {
        let b = Bytes::from(bytes.as_ref().to_vec());
        let bencoded = BencodeValue::decode(b.clone())?;
        let dict = bencoded.dict().ok_or(MetaInfoParsingError::WrongType)?;

        let info = dict
            .get(b"info")
            .ok_or(MetaInfoParsingError::MissingEntry)?
            .dict()
            .ok_or(MetaInfoParsingError::WrongType)?;
        let data = info.buffer();
        let info = MetaInfo::from_bytes(data)?;
        let name = dict
            .get(b"name")
            .and_then(|name| name.bytes())
            .map(|str| String::from_utf8(str.to_vec()))
            .transpose()?
            .or(name)
            .unwrap_or(info.info_hash.to_string());

        let mut announce_list = if let Some(announce) = dict.get(b"announce-list") {
            announce
                .list()
                .map(|list| -> Result<Vec<HashSet<Url>>, MetaInfoParsingError> {
                    let mut items = vec![];

                    for i in 0..list.len() {
                        let mut urls = HashSet::new();
                        let item = list.get(i).unwrap();

                        let url_list = item.list().unwrap();
                        for url_i in 0..url_list.len() {
                            urls.insert(
                                Url::parse(
                                    &*url_list.get(url_i).unwrap().string().transpose()?.unwrap(),
                                )
                                .unwrap(),
                            );
                        }

                        items.push(urls);
                    }

                    Ok(items)
                })
                .transpose()?
                .unwrap_or(vec![])
        } else {
            vec![]
        };

        if let Some(item) = dict.get(b"announce") {
            if let Some(Ok(url)) = item.string() {
                if let Ok(url) = Url::parse(&url) {
                    if announce_list.is_empty() {
                        announce_list.push(HashSet::new());
                    }

                    announce_list[0].insert(url);
                }
            }
        }

        Ok(Torrent {
            name,
            announce_list,
            info,
            bencoded,
        })
    }

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

    pub fn info_hash(&self) -> U160 {
        self.info.info_hash()
    }

    pub fn meta_info(&self) -> &MetaInfo {
        &self.info
    }

    pub fn announce_list(&self) -> &Vec<HashSet<Url>> {
        self.announce_list.as_ref()
    }
}

#[derive(Clone)]
pub struct MetaInfo {
    info_hash: U160,
    bencoded: BencodeValue,
    piece_length: usize,
    pieces: Bytes,
    object: MetaInfoObject,
}

impl MetaInfo {
    pub fn info_hash(&self) -> U160 {
        self.info_hash
    }

    pub fn piece_length(&self) -> usize {
        self.piece_length
    }

    pub fn pieces(&self) -> usize {
        self.pieces.len() / 20
    }

    pub fn hash(&self, index: usize) -> Bytes {
        self.pieces.slice(index * 20..(index + 1) * 20)
    }

    pub fn object(&self) -> &MetaInfoObject {
        &self.object
    }
}

#[derive(Clone, Debug)]
pub enum MetaInfoObject {
    Files(Vec<MetaFile>),
    File(usize),
}

impl Debug for MetaInfo {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MetaInfo")
            .field("info_hash", &self.info_hash)
            .field("piece_length", &self.piece_length)
            .field("object", &self.object)
            .finish()
    }
}

impl MetaInfo {
    pub fn from_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<Self, MetaInfoParsingError> {
        let b = Bytes::from(bytes.as_ref().to_vec());
        let bencoded = BencodeValue::decode(b)?;
        let dict = bencoded.dict().ok_or(MetaInfoParsingError::WrongType)?;

        let piece_length = dict
            .get(b"piece length")
            .ok_or(MetaInfoParsingError::MissingEntry)?
            .int()
            .ok_or(MetaInfoParsingError::WrongType)?
            .to_usize()
            .ok_or(MetaInfoParsingError::WrongType)?;

        let object = if let Some(files) = dict.get(b"files") {
            let mut meta_files = vec![];
            let files = files.list().ok_or(MetaInfoParsingError::WrongType)?;

            for file_i in 0..files.len() {
                let file = files.get(file_i).unwrap();
                meta_files.push(MetaFile::from_bencoded(file)?);
            }

            MetaInfoObject::Files(meta_files)
        } else {
            MetaInfoObject::File(
                dict.get(b"length")
                    .ok_or(MetaInfoParsingError::MissingEntry)?
                    .int()
                    .ok_or(MetaInfoParsingError::WrongType)?
                    .to_usize()
                    .ok_or(MetaInfoParsingError::WrongType)?,
            )
        };

        let pieces = dict
            .get(b"pieces")
            .ok_or(MetaInfoParsingError::MissingEntry)?
            .bytes()
            .ok_or(MetaInfoParsingError::WrongType)?;

        let info_hash =
            U160::from_bytes(digest(&SHA1_FOR_LEGACY_USE_ONLY, dict.buffer().as_ref())).unwrap();

        Ok(MetaInfo {
            info_hash,
            bencoded,
            piece_length,
            pieces,
            object,
        })
    }
}

#[derive(Debug, Clone)]
pub struct MetaFile {
    length: usize,
    path: Vec<String>,
}

impl MetaFile {
    fn from_bencoded(bencoded: &BencodeValue) -> Result<Self, MetaInfoParsingError> {
        let dict = bencoded.dict().ok_or(MetaInfoParsingError::WrongType)?;
        let path = dict
            .get(b"path")
            .ok_or(MetaInfoParsingError::MissingEntry)?
            .list()
            .ok_or(MetaInfoParsingError::WrongType)?;

        let mut path_ele = vec![];

        for path_i in 0..path.len() {
            let path_item = path.get(path_i).unwrap();
            path_ele.push(
                path_item
                    .string()
                    .ok_or(MetaInfoParsingError::WrongType)??,
            )
        }

        Ok(MetaFile {
            path: path_ele,
            length: dict
                .get(b"length")
                .ok_or(MetaInfoParsingError::MissingEntry)?
                .int()
                .ok_or(MetaInfoParsingError::WrongType)?
                .to_usize()
                .ok_or(MetaInfoParsingError::WrongType)?,
        })
    }

    pub fn path(&self) -> &[String] {
        &self.path
    }

    pub fn length(&self) -> usize {
        self.length
    }
}