Initial commit

src/main.rs

@@ -0,0 +1,269 @@
+extern crate cookie;
+extern crate dotenv;
+extern crate json;
+#[macro_use]
+extern crate lazy_static;
+extern crate reqwest;
+extern crate zip;
+
+use std::env;
+use std::fs::File;
+use std::io::{Read, BufReader, Write};
+use std::path::{Path, PathBuf};
+use std::collections::HashMap;
+use std::time::{Instant, Duration};
+
+use json::JsonValue;
+use reqwest::{Client, RedirectPolicy};
+use reqwest::header::{Cookie, Charset, SetCookie};
+use reqwest::header::{ContentDisposition, DispositionType, DispositionParam};
+use zip::ZipArchive;
+
+use catalogs::*;
+
+#[allow(dead_code)]
+mod catalogs {
+    type Login = Option<(&'static str, &'static str)>;
+    type CatalogDef = (&'static str, Login);
+
+    // (catalog_name, username, password)
+    pub const EIDI: CatalogDef = ("eidi1-2016", Some(("eidi-2016", "PGdP.16")));
+    pub const ERA: CatalogDef = ("era-2016", None);
+    pub const DS: CatalogDef = ("16w-diskrete-strukturen", None);
+}
+
+// TODO: use new Cookie release
+// TODO: use clap for proper CLI
+// TODO: clippy
+
+lazy_static! {
+    static ref CLIENT: Client = {
+        let mut client = Client::new().unwrap();
+        // The login site redirects to itself if no redirect parameter is given
+        client.redirect(RedirectPolicy::none());
+        client
+    };
+}
+
+#[derive(Debug)]
+struct Presentation {
+    name: String,
+    download_url: String,
+}
+
+impl<'a> From<&'a JsonValue> for Presentation {
+    fn from(json: &'a JsonValue) -> Presentation {
+        Presentation {
+            name: json["Name"].clone().take_string().unwrap(),
+            download_url: json["DownloadUrls"][0]["Url"].clone().take_string().unwrap(),
+        }
+    }
+}
+
+fn main() {
+    let catalog_def = DS;
+
+    let out_dir = get_output_directory();
+
+    let default_login = get_default_login();
+    let (catalog_name, (username, password)) = if let Some((user, pass)) = catalog_def.1 {
+        (catalog_def.0, (user.to_string(), pass.to_string()))
+    } else {
+        (catalog_def.0, default_login)
+    };
+
+    println!("Preparing to download catalog \"{}\"!", catalog_name);
+
+    let auth = get_auth(&username, &password);
+    // let out_dir = Path::new("/run/media/bobo1239/30E0E835E0E7FECA/downloaded");
+    // let out_dir = Path::new("out");
+    download_catalog(catalog_name, &out_dir, &auth);
+}
+
+fn get_output_directory() -> PathBuf {
+    // The first argument is (usually) the executable
+    let path_str = env::args_os().nth(1).expect("No output directory specified!");
+    let path = Path::new(&path_str);
+    if path.exists() {
+        assert!(path.is_dir());
+    } else {
+        ::std::fs::create_dir_all(path).unwrap();
+    }
+    path.to_owned()
+}
+
+fn get_default_login() -> (String, String) {
+    dotenv::dotenv().expect("Missing .env!");
+
+    let username = env::var("TUM_USERNAME").expect("Missing TUM_USERNAME environment variable!");
+    let password = env::var("TUM_PASSWORD").expect("Missing TUM_PASSWORD environment variable!");
+    (username, password)
+}
+
+fn get_auth(username: &str, password: &str) -> String {
+    println!("Logging in!");
+
+    let mut form_data = HashMap::new();
+    form_data.insert("UserName", username);
+    form_data.insert("Password", password);
+
+    let res = CLIENT.post("https://streams.tum.de/Mediasite/Login")
+        .form(&form_data)
+        .send()
+        .unwrap();
+
+    let set_cookie: &SetCookie = res.headers().get().unwrap();
+    let cookie = cookie::Cookie::parse(set_cookie.0[0].to_string()).unwrap();
+    assert_eq!(cookie.name(), "MediasiteAuth");
+
+    cookie.value().to_string()
+}
+
+fn download_catalog(catalog_name: &str, out_dir: &Path, auth: &str) {
+    let catalog_id = get_catalog_id(catalog_name, auth);
+    let json = get_json(&catalog_id, auth);
+    let presentations = json_to_presentations(&json);
+    println!("Starting to download {} presentations!",
+             presentations.len());
+
+    for presentation in presentations {
+        download_presentation(&presentation, out_dir, auth);
+    }
+}
+
+fn get_catalog_id(name: &str, auth: &str) -> String {
+    println!("Fetching catalog id!");
+
+    let url = format!("https://streams.tum.de/Mediasite/Catalog/catalogs/{}", name);
+    let res = CLIENT.get(&url)
+        .header(construct_cookie(auth))
+        .send();
+    let mut string = String::new();
+    res.unwrap().read_to_string(&mut string).unwrap();
+
+    let prefix = "CatalogId: '";
+    let idx = string.find(prefix).unwrap();
+    let pre_len = prefix.len();
+    // Assuming all catalog ids follow this pattern!
+    let len = "a6fca0c1-0be4-4e66-83b7-bcdc4eb5e95e".len();
+
+    string[(idx + pre_len)..(idx + pre_len + len)].to_string()
+}
+
+fn download_presentation(presentation: &Presentation, out_dir: &Path, auth: &str) {
+    println!("Downloading: {}", presentation.name);
+
+    let response = CLIENT.get(&presentation.download_url)
+        .header(construct_cookie(auth))
+        .send()
+        .unwrap();
+
+    let filename = {
+        let content_disposition: &ContentDisposition = response.headers().get().unwrap();
+        assert_eq!(content_disposition.disposition, DispositionType::Attachment);
+        let mut name = None;
+        for param in &content_disposition.parameters {
+            match *param {
+                DispositionParam::Filename(ref charset, _, ref vec) => {
+                    assert_eq!(&Charset::Ext("UTF-8".to_string()), charset);
+                    name = Some(String::from_utf8(vec.to_vec()).unwrap());
+                }
+                _ => continue,
+            }
+        }
+        name.expect("Missing filename in ContentDisposition!")
+    };
+
+    let mut path = PathBuf::from(out_dir);
+    path.push(fix_filename(&filename));
+    println!("Output file: {:?}", path);
+    if path.exists() {
+        if zip_is_valid(&path) {
+            println!("Already present and valid! Skipping!");
+            return;
+        } else {
+            println!("Found corrupt zip! Starting anew!");
+            // TODO?: We could try resuming the download if the server supports it
+            // http://stackoverflow.com/questions/41444297/
+            // http://stackoverflow.com/questions/3428102/
+        }
+    }
+    let mut file = File::create(&path).unwrap();
+
+    let mut reader = BufReader::new(response);
+    let mut buf = [0u8; 8 * 1024];
+
+    let update_duration_millis = 1000;
+    let mut last = Instant::now();
+    let mut done = 0;
+
+    loop {
+        let len = match reader.read(&mut buf) {
+            Ok(0) => break,
+            Ok(len) => len,
+            Err(ref e) if e.kind() == ::std::io::ErrorKind::Interrupted => continue,
+            Err(e) => panic!(e),
+        };
+
+        file.write_all(&buf[0..len]).unwrap();
+        done += len;
+
+        let now = Instant::now();
+        if now - last > Duration::from_millis(update_duration_millis) {
+            print!("{:.*} kB/s\r",
+                   2,
+                   done as f64 * 1000.0 / update_duration_millis as f64 / 1024.0);
+            std::io::stdout().flush().unwrap();
+            last = now;
+            done = 0;
+        }
+    }
+    ::std::mem::drop(file);
+
+    assert!(zip_is_valid(&path));
+}
+
+fn fix_filename(string: &str) -> String {
+    string.replace('/', "_")
+}
+
+fn json_to_presentations(json_str: &str) -> Vec<Presentation> {
+    let mut vec = Vec::new();
+    let mut json = json::parse(json_str).unwrap();
+    let mut count = 0;
+    for presentation in json["PresentationDetailsList"].members_mut() {
+        assert_eq!(1, presentation["DownloadUrls"].len());
+        vec.push(Presentation::from(&*presentation));
+        count += 1;
+    }
+    assert_eq!(count, json["TotalItems"]); // Maybe there are multiple pages
+    vec
+}
+
+fn get_json(catalog_id: &str, auth: &str) -> String {
+    println!("Fetching catalog!");
+
+    let mut data = HashMap::new();
+    data.insert("CatalogId", catalog_id);
+    data.insert("CurrentFolderId", catalog_id);
+    data.insert("ItemsPerPage", "500");
+
+    let res = CLIENT.post("https://streams.tum.de/Mediasite/Catalog/Data/GetPresentationsForFolder")
+        .header(construct_cookie(auth))
+        .json(&data)
+        .send();
+    let mut string = String::new();
+    res.unwrap().read_to_string(&mut string).unwrap();
+    string
+}
+
+fn construct_cookie(auth: &str) -> Cookie {
+    Cookie(vec![format!("MediasiteAuth={}", auth)])
+}
+
+fn zip_is_valid(path: &Path) -> bool {
+    assert!(path.exists());
+    assert!(path.is_file());
+    let file = File::open(path).unwrap();
+    ZipArchive::new(file).is_ok()
+}