bt/src/metainfo.rs

use std::path::PathBuf;

use sha1::Sha1;

use bencode::{encode, Object};

macro_rules! ts {
    ( $e:expr ) => {
        match $e {
            Some(x) => x,
            None => return None,
        }
    }
}

#[derive(Debug)]
pub struct Metainfo {
    pub announce: String,
    pub announce_list: Vec<Vec<String>>,
    pub files: Vec<MetainfoFile>,
    pub info_hash: [u8; 20],
    pub piece_length: u64,
    pub pieces: Vec<[u8; 20]>,
}

impl Metainfo {
    pub fn from_bencode<O>(obj: O) -> Option<Metainfo> where O: Into<Object> {
        let obj = obj.into();
        let metainfo = ts!(obj.as_dict());
        let info = ts!(metainfo.get_dict("info"));

        let announce = ts!(metainfo.get_string("announce"));
        let info_hash = sha1(&encode(ts!(metainfo.get_object("info"))));
        let name = ts!(info.get_str("name"));
        let piece_length = ts!(info.get_int("piece length"));
        let pieces = ts!(info.get_bytes("pieces"));

        let mut announce_list = vec![];
        if let Some(list) = metainfo.get_list("announce-list") {
            for tier in list.iter() {
                let mut urls = vec![];
                for url in ts!(tier.as_list()).iter() {
                    urls.push(ts!(url.as_string()));
                }
                announce_list.push(urls);
            }
        }

        let mut files = vec![];
        if let Some(list) = info.get_list("files") {
            for obj in list.iter() {
                let file = ts!(obj.as_dict());

                let mut path = PathBuf::new();
                path.push(name);
                for component in ts!(file.get_list("path")).iter() {
                    path.push(ts!(component.as_str()));
                }

                files.push(MetainfoFile {
                    length: ts!(file.get_int("length")) as u64,
                    path: path,
                });
            }
        } else {
            let mut path = PathBuf::new();
            path.push(name);

            files.push(MetainfoFile{
                length: ts!(info.get_int("length")) as u64,
                path: path,
            })
        }

        Some(Metainfo {
            announce: announce,
            announce_list: announce_list,
            files: files,
            info_hash: info_hash,
            piece_length: piece_length as u64,
            pieces: pieces.chunks(20).map(|c| piece_from_slice(c)).collect::<Vec<_>>(),
        })
    }
}

#[derive(Debug)]
pub struct MetainfoFile {
    pub length: u64,
    pub path: PathBuf,
}

fn sha1(bytes: &[u8]) -> [u8; 20] {
    let mut hasher = Sha1::new();
    hasher.update(bytes);
    hasher.digest().bytes()
}

fn piece_from_slice(src: &[u8]) -> [u8; 20] {
    assert_eq!(src.len(), 20);

    let mut dst = [0u8; 20];
    dst.copy_from_slice(src);
    dst
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use bencode::{Dict, List, Object};
    use metainfo::Metainfo;

    #[test]
    fn test_single_file() {
        let mut meta = Dict::new();
        meta.insert("announce", "http://ubuntu.com/tracker:6969");

        let mut info = Dict::new();
        info.insert("name", "ubuntu-16.04-desktop.iso");
        info.insert("length", 1024 * 1024 * 1024);
        info.insert("piece length", 1024 * 512);
        info.insert("pieces", "");

        meta.insert("info", info);

        let metainfo = Metainfo::from_bencode(meta).unwrap();
        assert_eq!(metainfo.announce, "http://ubuntu.com/tracker:6969");
        assert_eq!(metainfo.piece_length, 1024 * 512);

        let file = &metainfo.files[0];
        assert_eq!(file.path.to_str().unwrap(), "ubuntu-16.04-desktop.iso");
        assert_eq!(file.length, 1024 * 1024 * 1024);
    }

    #[test]
    fn test_multiple_files() {
        let mut meta = Dict::new();
        meta.insert("announce", "http://ubuntu.com/tracker:6969");

        let mut info = Dict::new();
        info.insert("name", "base_folder");
        info.insert("piece length", 1024 * 512);
        info.insert("pieces", "");

        let mut files = List::new();
        let mut file1 = Dict::new();
        file1.insert("path", vec!["folder2", "image.txt"]);
        file1.insert("length", 512);
        files.push(file1);
        let mut file2 = Dict::new();
        file2.insert("path", vec!["txt.jpg"]);
        file2.insert("length", 64);
        files.push(file2);
        info.insert("files", Object::List(files));

        meta.insert("info", info);

        let metainfo = Metainfo::from_bencode(meta).unwrap();
        assert_eq!(metainfo.announce, "http://ubuntu.com/tracker:6969");
        assert_eq!(metainfo.piece_length, 1024 * 512);

        let file1 = &metainfo.files[0];
        assert_eq!(file1.path, Path::new("base_folder/folder2/image.txt"));
        assert_eq!(file1.length, 512);

        let file2 = &metainfo.files[1];
        assert_eq!(file2.path, Path::new("base_folder/txt.jpg"));
        assert_eq!(file2.length, 64);
    }
}