torrent-rs/src/main.rs

use std::{env, fs, path};
use std::collections::HashMap;

use anyhow::{anyhow, Result};
use rand::prelude::*;
use reqwest::blocking::Client;
use reqwest::Url;

use crate::bencode::{Bencode, ByteString};

mod bencode;

#[derive(Debug)]
struct FileInfo {
    length: i64,
    path: String,
}

#[derive(Debug)]
struct TorrentInfo {
    files: Option<Vec<FileInfo>>,
    length: Option<i64>,
    name: String,
    piece_length: i64,
    pieces: Vec<[u8; 20]>,
}

#[derive(Debug)]
struct Torrent {
    announce: String,
    info: TorrentInfo,
    info_hash: [u8; 20],
    _extra: HashMap<String, String>
}

impl Torrent {
    pub fn from(input: &Bencode) -> Result<Self> {
        let Bencode::Dict(decoded) = input else {
            return Err(anyhow!("Torrent has wrong type. Not a dict"))
        };
        let b_info = decoded.get(&ByteString::from_str("info"))
            .ok_or(anyhow!("Could not find key \"info\""))?;
        let Bencode::Dict(info_dict) = b_info else {
            return Err(anyhow!("Info has wrong type. Not a dict"))
        };

        let Bencode::Bytes(name) = info_dict.get(&ByteString::from_str("name"))
            .ok_or(anyhow!("Could not find key \"name\""))? else {
            return Err(anyhow!("name not a string"))
        };

        let Bencode::Integer(piece_length) = *info_dict.get(&ByteString::from_str("piece length"))
            .ok_or(anyhow!("Could not find key \"piece length\""))? else {
            return Err(anyhow!("piece length not an integer"))
        };

        let Bencode::Bytes(all_pieces) = info_dict.get(&ByteString::from_str("pieces"))
            .ok_or(anyhow!("Could not find key \"pieces\""))? else {
            return Err(anyhow!("pieces not an string"))
        };

        if all_pieces.to_vec().len() % 20 != 0 {
            return Err(anyhow!("Pieces string length not a multiple of 20"))
        }

        // TODO maybe handle the error, even though it should be checked above
        let pieces = all_pieces.to_vec().chunks(20).map(|x| x.try_into().unwrap()).collect();

        let info;
        if let Some(b_length) = info_dict.get(&ByteString::from_str("length")) {
            let Bencode::Integer(length) = b_length else {
                return Err(anyhow!("length not an integer"))
            };
            info = TorrentInfo {
                files: None,
                length: Some(*length),
                name: name.to_string()?,
                piece_length,
                pieces,
            }
        } else if let Some(b_files) = info_dict.get(&ByteString::from_str("files")) {
            let Bencode::List(files_list) = b_files else {
                return Err(anyhow!("files not a list"))
            };

            let mut files = Vec::new();
            for b_file_info in files_list {
                let Bencode::Dict(file_info) = b_file_info else {
                    return Err(anyhow!("file_info not a dict"))
                };
                let Bencode::List(b_path) = file_info.get(&ByteString::from_str("path"))
                    .ok_or(anyhow!("Could not find key \"path\""))? else {
                    return Err(anyhow!("path not a list"))
                };
                let mut v_path = Vec::new();
                for b_path_part in b_path {
                    let Bencode::Bytes(path_part) = b_path_part else {
                        return Err(anyhow!("path part not a string"))
                    };
                    v_path.push(path_part.to_string()?);
                }

                let Bencode::Integer(length) = *file_info.get(&ByteString::from_str("length"))
                    .ok_or(anyhow!("Could not find key \"length\""))? else {
                    return Err(anyhow!("length not an integer"))
                };
                files.push(FileInfo{
                    path: v_path.join(path::MAIN_SEPARATOR_STR),
                    length
                })
            }

            info = TorrentInfo {
                files: Some(files),
                length: None,
                name: name.to_string()?,
                piece_length,
                pieces,
            }
        } else {
            return Err(anyhow!("Could not find key \"length\" or \"files\""))
        }


        let b_announce = decoded.get(&ByteString::from_str("announce"))
            .ok_or(anyhow!("Could not find key \"announce\""))?;
        let Bencode::Bytes(announce) = b_announce else {
            return Err(anyhow!("announce not a string"))
        };

        let mut _extra = HashMap::new();
        for (k, v) in decoded {
            if k == &ByteString::from_str("announce") || k == &ByteString::from_str("info") {
                continue
            }
            let Ok(k_str) = k.to_string() else {
                println!("Ignoring non-UTF-8 key {k:?}");
                continue;
            };
            match v {
                Bencode::Integer(i) => {
                    _extra.insert(k_str, format!("{}", i));
                }
                Bencode::Bytes(s) => {
                    let Ok(s_str) = s.to_string() else {
                        println!("Ignoring non-UTF-8 value of key {k_str}");
                        continue;
                    };
                    _extra.insert(k_str, s_str);
                }
                _ => {
                    println!("Ignoring extra of type list or dict")
                }
            }
        }

        let result = Self {
            announce: announce.to_string()?,
            info_hash: b_info.sha1(),
            info,
            _extra,
        };
        Ok(result)
    }
}

fn main() -> Result<()> {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        return Err(anyhow!("Please specify the torrent file to download"));
    }
    let file_path = &args[1];

    let torrent_bytes = fs::read(&file_path)?;

    let torrent_parsed = Bencode::parse(&torrent_bytes)?;

    println!("torrent file parsed");
    // println!("parsed file: {torrent_parsed}");

    let torrent = Torrent::from(&torrent_parsed)?;

    // println!("torrent decoded: {torrent:?}");

    let torrent_length= torrent.info.length.unwrap_or(
        torrent.info.files.unwrap().iter().map(|x| x.length).sum()
    );


    let mut peer_id= [0u8; 20];
    thread_rng().fill_bytes(&mut peer_id);

    println!("Generated peer id: {}",
             peer_id.iter().map(|x| {format!("{x:02x}")}).collect::<Vec<_>>().join(""));
    println!("Making request to tracker");

    println!("info hash: {}", torrent.info_hash.iter().map(|x| {format!("{x:02x}")}).collect::<Vec<_>>().join(""));

    let mut url = Url::parse(&torrent.announce)?;
    url.query_pairs_mut()
        .append_pair("port", "6885")
        .append_pair("uploaded", "0")
        .append_pair("downloaded", "0")
        .append_pair("compact", "1")
        .append_pair("left", &format!("{torrent_length}"));


    let mut query = url.query().unwrap().to_string();
    query = format!("{query}&peer_id={}&info_hash={}",
                     urlencoding::encode_binary(&peer_id[..]),
                     urlencoding::encode_binary(&torrent.info_hash[..]));
    url.set_query(Some(&query));

    println!("url: {}", url);

    let client = Client::new();
    let resp = client.get(url).send()?;
    let status = resp.status();

    let body = resp.text()?;
    println!("Response: {} {}", status, body);

    Ok(())
}