<commit_before>use std::io::prelude::*;
use std::{process, env};
use std::path::Path;
use std::error::Error;
use std::fs::{self, OpenOptions};
/* Alias to avoid naming conflict for write_all!() */
use std::fmt::Write as std_write;

use url::Url;
use url_open::UrlOpen;

use serde_json::{from_str, to_string};

use util::*;
use util::{git_author, git_commit_message, format_file};
use timesheet::traits::HasHTML;

use timesheet::session::Session;
use timesheet::session::EventType;

#[derive(Serialize, Deserialize, Debug)]
pub struct Timesheet {
    start            : u64,
    end              : u64,
    user             : String,
    pub show_commits : bool,
    repo             : String,
    sessions         : Vec<Session>,
}

impl Timesheet {
    // TODO: check if i can just write_files before the end of main()
    /** Initializes the .trk/timesheet.json file which holds
     * the serialized timesheet
     * Returns Some(newTimesheet) if operation succeeded */
    pub fn init(author_name: Option<&str>) -> Option<Timesheet> {
        /* Check if file already exists (no init permitted) */
        if Timesheet::is_init() {
            println!("Timesheet is already initialized!");
            return None;
        }

        /* File does not exist, initialize */
        let git_author_name = git_author();
        let author_name = match author_name {
            Some(name) => name,
            None => {
                match git_author_name {
                    Some(ref git_name) => git_name,
                    None => {
                        println!("Empty name not permitted.
Please run with 'trk init <name>'");
                        process::exit(0);
                    }
                }
            }
        };
        let now = get_seconds();
        let sheet = Timesheet {
            start        : now,
            end          : now + 1,
            user         : author_name.to_string(),
            show_commits : true,
            repo         : String::new(),
            sessions     : Vec::<Session>::new(),
        };
        if sheet.write_files() {
            Some(sheet)
        } else {
            None // TODO: error message?
        }
    }

    fn is_init() -> bool {
        Path::new("./.trk/timesheet.json").exists() && Timesheet::load_from_file().is_some()
    }

    pub fn new_session(&mut self, timestamp: Option<u64>) -> bool {
        let possible = match self.sessions.last_mut() {
            None => true,
            Some(session) => {
                if session.is_running() {
                    println!("Last session is still running.");
                    false
                } else {
                    true
                }
            }
        };
        if possible {
            match timestamp {
                Some(timestamp) => {
                    let is_valid_ts = match self.sessions.last() {
                        None => timestamp > self.start,
                        Some(last_session) => timestamp > last_session.end,
                    };
                    if is_valid_ts {
                        self.sessions.push(Session::new(Some(timestamp)));
                    } else {
                        println!("That timestamp is invalid.");
                        process::exit(0);
                    }
                }
                None => {
                    self.sessions.push(Session::new(None));
                }
            };
            self.write_files();
        }
        possible
    }

    pub fn end_session(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
		// TODO This is always problematic - rethink.
                session.update_end();
                session.finalize(timestamp);
                self.end = session.end + 1;
            }
            None => println!("No session to finalize."),
        }
        self.write_files();
    }

    pub fn pause(&mut self, timestamp: Option<u64>, note: Option<String>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, note, EventType::Pause);
            }
            None => println!("No session to pause."),
        }
        self.write_files();
    }

    pub fn resume(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, None, EventType::Resume);
            }
            None => println!("No session to pause."),
        }
        self.write_files();
    }

    pub fn note(&mut self, timestamp: Option<u64>, note_text: String) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, Some(note_text), EventType::Note);
            }
            None => println!("No session to add note to."),
        }
        self.write_files();
    }

    pub fn add_commit(&mut self, hash: String) {
        let new_needed = match self.sessions.last() {
            Some(session) => !session.is_running(),
            None => true,
        };
        if new_needed {
            self.new_session(None);
            self.write_files();
        }
        match self.sessions.last_mut() {
            Some(session) => {
                let message = git_commit_message(&hash).unwrap_or(String::new());
                session.push_event(None, Some(message), EventType::Commit { hash });
            }
            None => println!("No session to add commit to."),
        }
        self.write_files();
    }

    pub fn add_branch(&mut self, name: String) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.add_branch(name);
            }
            None => {}
        }
        self.write_files();
    }

    pub fn write_to_html(&self, ago: Option<u64>) -> bool {
        /* TODO: avoid time-of-check-to-time-of-use race risk */
        /* TODO: make all commands run regardless of where trk is executed
         * (and not just in root which is assumed here */

        let path = Path::new("./timesheet.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                file.write_all(self.to_html(ago).as_bytes()).unwrap();
                format_file("timesheet.html");
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not report sheet! {}", why.description());
                false
            }
        }
    }

    pub fn write_last_session_html(&self) -> bool {
        let path = Path::new("./session.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        let mut file = match file {
            Ok(file) => file,
            Err(why) => {
                println!("Could not write report to session.html! {}",
                        why.description());
                return false;
            }
        };

        let session = match self.sessions.last() {
            Some(session) => session,
            /* TODO: write empty file anyway? */
            None => return true,
        };

        let stylesheets = match self.show_commits {
            true  => r#"<link rel="stylesheet" type="text/css" href="style.css">
"#,
            false => r#"<link rel="stylesheet" type="text/css" href="style.css">
<link rel="stylesheet" type="text/css" href="no_commit.css">
"#,
        };

        let html = format!(r#"<!DOCTYPE html>

  {}
  <title>{} for {}</title>

{}

"#,
                   stylesheets,
                   "Session",
                   "Rafael Bachmann",
                   session.to_html());
    file.write_all(html.as_bytes()).unwrap();
    format_file("session.html");
    /* Save was successful */
    true
    }

    fn write_to_json(&self) -> bool {
        if !Path::new("./.trk").exists() {
            match fs::create_dir("./.trk") {
                Ok(_) => {}
                _ => {
                    println!("Could not create .trk directory.");
                    process::exit(0);
                }
            }
        }

        let path = Path::new("./.trk/timesheet.json");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                /* Convert the sheet to a JSON string. */
                let serialized =
                    to_string(&self).expect("Could not write serialized time sheet.");
                file.write_all(serialized.as_bytes()).unwrap();
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not open timesheet.json file: {}", why.description());
                false
            }
        }
    }

    fn write_files(&self) -> bool {
        /* TODO: avoid time-of-check-to-time-of-use race risk */
        /* TODO: make all commands run regardless of where trk is executed
         * (and not just in root which is assumed here */
        self.write_to_json() && self.write_to_html(None) && self.write_last_session_html()
    }

    /** Return a Some(Timesheet) struct if a timesheet.json file
     * is present and valid in the .trk directory, and None otherwise.
     * TODO: improve error handling
     * */
    pub fn load_from_file() -> Option<Timesheet> {
        let path = Path::new("./.trk/timesheet.json");
        let file = OpenOptions::new().read(true).open(&path);
        match file {
            Ok(mut file) => {
                let mut serialized = String::new();
                match file.read_to_string(&mut serialized) {
                    Ok(..) => from_str(&serialized).unwrap_or(None),
                    Err(..) => {
                        println!("IO error while reading the timesheet file.");
                        process::exit(0);
                    }
                }
            }
            Err(..) => None,
        }
    }

    /* TODO: Check why timestamps aren't overwritten */
    pub fn clear() {
        /* Try to get user name */
        let sheet = Timesheet::load_from_file();
        let name: Option<String> = sheet.map(|s| s.user.clone());

        let path = Path::new("./.trk/timesheet.json");
        if path.exists() {
            match fs::remove_file(&path) {
                Ok(..) => {}
                Err(why) => println!("Could not remove sessions file: {}", why.description()),
            }
        }
        match name {
            Some(name) => {
                /* Overwrite file */
                Timesheet::init(Some(&name));
            }
            None => {
                Timesheet::init(None);
            }
        }
    }

    pub fn timesheet_status(&self) -> String {
        let mut status = format!("Sheet running for {}\n",
               sec_to_hms_string(get_seconds() - self.start));
        match self.sessions.len() {
            0 => write!(&mut status, "No sessions yet.\n").unwrap(),
            n => {
                write!(&mut status,
                       "{} session(s) so far.\nLast session:\n{}",
                       n,
                       self.sessions[n - 1].status())
                        .unwrap()
            }
        };
        status
    }

    pub fn last_session_status(&self) -> String {
        let status = self.sessions.last().map(|session| session.status());
        match status {
            None => "No session yet.".to_string(),
            Some(status) => status,
        }
    }

    pub fn report_last_session(&self) {
        // We assume that we are in a valid directory.
        let mut p = env::current_dir().unwrap();
        p.push("session.html");
        let path = p.as_path();
        let mut path_str = "file://".to_string();
        path_str.push_str(path.to_str().unwrap());
        Url::parse(&path_str).unwrap().open();
    }

    pub fn report_sheet(&self, ago: Option<u64>) {
        self.write_to_html(ago);
        // We assume that we are in a valid directory.
        let mut p = env::current_dir().unwrap();
        p.push("timesheet.html");
        let path = p.as_path();
        let mut path_str = "file://".to_string();
        path_str.push_str(path.to_str().unwrap());
        Url::parse(&path_str).unwrap().open();
        /* clean up html again */
        self.write_to_html(None);
    }

    pub fn toggle_show_git_info(&mut self, setting: bool) {
        self.show_commits = setting;
        self.write_files();
    }

    pub fn set_repo_url(&mut self, repo: String) {
        self.repo = repo;
    }

    pub fn pause_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.pause_time())
    }

    pub fn work_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.work_time())
    }

    fn to_html(&self, ago: Option<u64>) -> String {
        let timestamp = match ago {
            Some(ago) => ago,
            None      => self.start,
        };
        let mut sessions_html = String::new();
        for session in &self.sessions {
            if session.start > timestamp {
		write!(&mut sessions_html, "{}
", session.to_html()).unwrap();
            }
        }

        let stylesheets = match self.show_commits {
            true => {
                r#"<link rel="stylesheet" type="text/css" href="style.css">
"#
                        .to_string()
            }
            false => {
                r#"
<link rel="stylesheet" type="text/css" href="style.css">
<link rel="stylesheet" type="text/css" href="no_commit.css">
"#
                        .to_string()
            }
        };

        let mut html = format!(r#"<!DOCTYPE html>

        {}
        <title>{} for {}</title>

    {}"#,
                               stylesheets,
                               "Timesheet",
                               "Rafael Bachmann",
                               sessions_html);

        write!(&mut html,
               r#"

Worked for {}

Paused for {}

"#,
               sec_to_hms_string(self.work_time()),
               sec_to_hms_string(self.pause_time()))
                .unwrap();
        write!(&mut html, "\n").unwrap();
        html
    }
}
<commit_msg>shorten to_html() for timesheet<commit_after>use std::io::prelude::*;
use std::{process, env};
use std::path::Path;
use std::error::Error;
use std::fs::{self, OpenOptions};
/* Alias to avoid naming conflict for write_all!() */
use std::fmt::Write as std_write;

use url::Url;
use url_open::UrlOpen;

use serde_json::{from_str, to_string};

use util::*;
use util::{git_author, git_commit_message, format_file};
use timesheet::traits::HasHTML;

use timesheet::session::Session;
use timesheet::session::EventType;

#[derive(Serialize, Deserialize, Debug)]
pub struct Timesheet {
    start            : u64,
    end              : u64,
    user             : String,
    pub show_commits : bool,
    repo             : String,
    sessions         : Vec<Session>,
}

impl Timesheet {
    // TODO: check if i can just write_files before the end of main()
    /** Initializes the .trk/timesheet.json file which holds
     * the serialized timesheet
     * Returns Some(newTimesheet) if operation succeeded */
    pub fn init(author_name: Option<&str>) -> Option<Timesheet> {
        /* Check if file already exists (no init permitted) */
        if Timesheet::is_init() {
            println!("Timesheet is already initialized!");
            return None;
        }

        /* File does not exist, initialize */
        let git_author_name = git_author();
        let author_name = match author_name {
            Some(name) => name,
            None => {
                match git_author_name {
                    Some(ref git_name) => git_name,
                    None => {
                        println!("Empty name not permitted.
Please run with 'trk init <name>'");
                        process::exit(0);
                    }
                }
            }
        };
        let now = get_seconds();
        let sheet = Timesheet {
            start        : now,
            end          : now + 1,
            user         : author_name.to_string(),
            show_commits : true,
            repo         : String::new(),
            sessions     : Vec::<Session>::new(),
        };
        if sheet.write_files() {
            Some(sheet)
        } else {
            None // TODO: error message?
        }
    }

    fn is_init() -> bool {
        Path::new("./.trk/timesheet.json").exists() && Timesheet::load_from_file().is_some()
    }

    pub fn new_session(&mut self, timestamp: Option<u64>) -> bool {
        let possible = match self.sessions.last_mut() {
            None => true,
            Some(session) => {
                if session.is_running() {
                    println!("Last session is still running.");
                    false
                } else {
                    true
                }
            }
        };
        if possible {
            match timestamp {
                Some(timestamp) => {
                    let is_valid_ts = match self.sessions.last() {
                        None => timestamp > self.start,
                        Some(last_session) => timestamp > last_session.end,
                    };
                    if is_valid_ts {
                        self.sessions.push(Session::new(Some(timestamp)));
                    } else {
                        println!("That timestamp is invalid.");
                        process::exit(0);
                    }
                }
                None => {
                    self.sessions.push(Session::new(None));
                }
            };
            self.write_files();
        }
        possible
    }

    pub fn end_session(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
		// TODO This is always problematic - rethink.
                session.update_end();
                session.finalize(timestamp);
                self.end = session.end + 1;
            }
            None => println!("No session to finalize."),
        }
        self.write_files();
    }

    pub fn pause(&mut self, timestamp: Option<u64>, note: Option<String>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, note, EventType::Pause);
            }
            None => println!("No session to pause."),
        }
        self.write_files();
    }

    pub fn resume(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, None, EventType::Resume);
            }
            None => println!("No session to pause."),
        }
        self.write_files();
    }

    pub fn note(&mut self, timestamp: Option<u64>, note_text: String) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, Some(note_text), EventType::Note);
            }
            None => println!("No session to add note to."),
        }
        self.write_files();
    }

    pub fn add_commit(&mut self, hash: String) {
        let new_needed = match self.sessions.last() {
            Some(session) => !session.is_running(),
            None => true,
        };
        if new_needed {
            self.new_session(None);
            self.write_files();
        }
        match self.sessions.last_mut() {
            Some(session) => {
                let message = git_commit_message(&hash).unwrap_or(String::new());
                session.push_event(None, Some(message), EventType::Commit { hash });
            }
            None => println!("No session to add commit to."),
        }
        self.write_files();
    }

    pub fn add_branch(&mut self, name: String) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.add_branch(name);
            }
            None => {}
        }
        self.write_files();
    }

    pub fn write_to_html(&self, ago: Option<u64>) -> bool {
        /* TODO: avoid time-of-check-to-time-of-use race risk */
        /* TODO: make all commands run regardless of where trk is executed
         * (and not just in root which is assumed here */

        let path = Path::new("./timesheet.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                file.write_all(self.to_html(ago).as_bytes()).unwrap();
                format_file("timesheet.html");
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not report sheet! {}", why.description());
                false
            }
        }
    }

    pub fn write_last_session_html(&self) -> bool {
        let path = Path::new("./session.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        let mut file = match file {
            Ok(file) => file,
            Err(why) => {
                println!("Could not write report to session.html! {}",
                        why.description());
                return false;
            }
        };

        let session = match self.sessions.last() {
            Some(session) => session,
            /* TODO: write empty file anyway? */
            None => return true,
        };

        let stylesheets = match self.show_commits {
            true  => r#"<link rel="stylesheet" type="text/css" href="style.css">
"#,
            false => r#"<link rel="stylesheet" type="text/css" href="style.css">
<link rel="stylesheet" type="text/css" href="no_commit.css">
"#,
        };

        let html = format!(r#"<!DOCTYPE html>

  {}
  <title>{} for {}</title>

{}

"#,
                   stylesheets,
                   "Session",
                   "Rafael Bachmann",
                   session.to_html());
    file.write_all(html.as_bytes()).unwrap();
    format_file("session.html");
    /* Save was successful */
    true
    }

    fn write_to_json(&self) -> bool {
        if !Path::new("./.trk").exists() {
            match fs::create_dir("./.trk") {
                Ok(_) => {}
                _ => {
                    println!("Could not create .trk directory.");
                    process::exit(0);
                }
            }
        }

        let path = Path::new("./.trk/timesheet.json");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                /* Convert the sheet to a JSON string. */
                let serialized =
                    to_string(&self).expect("Could not write serialized time sheet.");
                file.write_all(serialized.as_bytes()).unwrap();
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not open timesheet.json file: {}", why.description());
                false
            }
        }
    }

    fn write_files(&self) -> bool {
        /* TODO: avoid time-of-check-to-time-of-use race risk */
        /* TODO: make all commands run regardless of where trk is executed
         * (and not just in root which is assumed here */
        self.write_to_json() && self.write_to_html(None) && self.write_last_session_html()
    }

    /** Return a Some(Timesheet) struct if a timesheet.json file
     * is present and valid in the .trk directory, and None otherwise.
     * TODO: improve error handling
     * */
    pub fn load_from_file() -> Option<Timesheet> {
        let path = Path::new("./.trk/timesheet.json");
        let file = OpenOptions::new().read(true).open(&path);
        match file {
            Ok(mut file) => {
                let mut serialized = String::new();
                match file.read_to_string(&mut serialized) {
                    Ok(..) => from_str(&serialized).unwrap_or(None),
                    Err(..) => {
                        println!("IO error while reading the timesheet file.");
                        process::exit(0);
                    }
                }
            }
            Err(..) => None,
        }
    }

    /* TODO: Check why timestamps aren't overwritten */
    pub fn clear() {
        /* Try to get user name */
        let sheet = Timesheet::load_from_file();
        let name: Option<String> = sheet.map(|s| s.user.clone());

        let path = Path::new("./.trk/timesheet.json");
        if path.exists() {
            match fs::remove_file(&path) {
                Ok(..) => {}
                Err(why) => println!("Could not remove sessions file: {}", why.description()),
            }
        }
        match name {
            Some(name) => {
                /* Overwrite file */
                Timesheet::init(Some(&name));
            }
            None => {
                Timesheet::init(None);
            }
        }
    }

    pub fn timesheet_status(&self) -> String {
        let mut status = format!("Sheet running for {}\n",
               sec_to_hms_string(get_seconds() - self.start));
        match self.sessions.len() {
            0 => write!(&mut status, "No sessions yet.\n").unwrap(),
            n => {
                write!(&mut status,
                       "{} session(s) so far.\nLast session:\n{}",
                       n,
                       self.sessions[n - 1].status())
                        .unwrap()
            }
        };
        status
    }

    pub fn last_session_status(&self) -> String {
        let status = self.sessions.last().map(|session| session.status());
        match status {
            None => "No session yet.".to_string(),
            Some(status) => status,
        }
    }

    pub fn report_last_session(&self) {
        // We assume that we are in a valid directory.
        let mut p = env::current_dir().unwrap();
        p.push("session.html");
        let path = p.as_path();
        let mut path_str = "file://".to_string();
        path_str.push_str(path.to_str().unwrap());
        Url::parse(&path_str).unwrap().open();
    }

    pub fn report_sheet(&self, ago: Option<u64>) {
        self.write_to_html(ago);
        // We assume that we are in a valid directory.
        let mut p = env::current_dir().unwrap();
        p.push("timesheet.html");
        let path = p.as_path();
        let mut path_str = "file://".to_string();
        path_str.push_str(path.to_str().unwrap());
        Url::parse(&path_str).unwrap().open();
        /* clean up html again */
        self.write_to_html(None);
    }

    pub fn toggle_show_git_info(&mut self, setting: bool) {
        self.show_commits = setting;
        self.write_files();
    }

    pub fn set_repo_url(&mut self, repo: String) {
        self.repo = repo;
    }

    pub fn pause_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.pause_time())
    }

    pub fn work_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.work_time())
    }

    fn to_html(&self, ago: Option<u64>) -> String {
        let timestamp = ago.unwrap_or(self.start);
        let mut sessions_html = String::new();
        for session in &self.sessions {
            if session.start > timestamp {
		write!(&mut sessions_html, "{}
", session.to_html()).unwrap();
            }
        }

        let stylesheets = if self.show_commits {
                "<link rel=\"stylesheet\" type=\"text/css\" href=\"style.css\">\n".to_string()
        } else {
                r#"<link rel="stylesheet" type="text/css" href="style.css">
<link rel="stylesheet" type="text/css" href="no_commit.css">
"#.to_string()
        };

        let mut html = format!(r#"<!DOCTYPE html>

        {}
        <title>{} for {}</title>

    {}"#,
              stylesheets,
              "Timesheet",
              "Rafael Bachmann",
              sessions_html);

        write!(&mut html,
               r#"

Worked for {}

Paused for {}

"#,
               sec_to_hms_string(self.work_time()),
               sec_to_hms_string(self.pause_time()))
               .unwrap();
        write!(&mut html, "\n").unwrap();
        html
    }
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::{self, Write};

use color;
use image::{ImageEncoder, ImageError, ImageResult};

const BITMAPFILEHEADER_SIZE: u32 = 14;
const BITMAPINFOHEADER_SIZE: u32 = 40;
const BITMAPV4HEADER_SIZE: u32 = 108;

/// The representation of a BMP encoder.
pub struct BMPEncoder<'a, W: 'a> {
    writer: &'a mut W,
}

impl<'a, W: Write + 'a> BMPEncoder<'a, W> {
    /// Create a new encoder that writes its output to ```w```.
    pub fn new(w: &'a mut W) -> Self {
        BMPEncoder { writer: w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```.
    pub fn encode(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
    ) -> ImageResult<()> {
        let bmp_header_size = BITMAPFILEHEADER_SIZE;

        let (dib_header_size, written_pixel_size, palette_color_count) = get_pixel_info(c)?;
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes
        let image_size = width.checked_mul(height);
        if image_size.is_none() {
            return Err(ImageError::DimensionError);
        }
        let image_size = image_size.unwrap().checked_mul(written_pixel_size);
        if image_size.is_none() {
            return Err(ImageError::DimensionError);
        }
        let image_size = image_size.unwrap() + (height * row_pad_size);
        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

        // write BMP header
        self.writer.write_u8(b'B')?;
        self.writer.write_u8(b'M')?;
        self.writer.write_u32::<LittleEndian>(file_size)?; // file size
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 1
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 2
        self.writer
            .write_u32::<LittleEndian>(bmp_header_size + dib_header_size + palette_size)?; // image data offset

        // write DIB header
        self.writer.write_u32::<LittleEndian>(dib_header_size)?;
        self.writer.write_i32::<LittleEndian>(width as i32)?;
        self.writer.write_i32::<LittleEndian>(height as i32)?;
        self.writer.write_u16::<LittleEndian>(1)?; // color planes
        self.writer
            .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; // bits per pixel
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(3)?; // compression method - bitfields
        } else {
            self.writer.write_u32::<LittleEndian>(0)?; // compression method - no compression
        }
        self.writer.write_u32::<LittleEndian>(image_size)?;
        self.writer.write_i32::<LittleEndian>(0)?; // horizontal ppm
        self.writer.write_i32::<LittleEndian>(0)?; // vertical ppm
        self.writer.write_u32::<LittleEndian>(palette_color_count)?;
        self.writer.write_u32::<LittleEndian>(0)?; // all colors are important
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(0xff << 16)?; // red mask
            self.writer.write_u32::<LittleEndian>(0xff << 8)?; // green mask
            self.writer.write_u32::<LittleEndian>(0xff << 0)?; // blue mask
            self.writer.write_u32::<LittleEndian>(0xff << 24)?; // alpha mask
            self.writer.write_u32::<LittleEndian>(0x73524742)?; // colorspace - sRGB

            // endpoints (3x3) and gamma (3)
            for _ in 0..12 {
                self.writer.write_u32::<LittleEndian>(0)?;
            }
        }

        // write image data
        match c {
            color::ColorType::Rgb8 => {
                self.encode_rgb(image, width, height, row_pad_size, 3)?
            }
            color::ColorType::Rgba8 => {
                self.encode_rgba(image, width, height, row_pad_size, 4)?
            }
            color::ColorType::L8 => {
                self.encode_gray(image, width, height, row_pad_size, 1)?
            }
            color::ColorType::La8 => {
                self.encode_gray(image, width, height, row_pad_size, 2)?
            }
            _ => {
                return Err(ImageError::IoError(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                )))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                let r = image[pixel_start];
                let g = image[pixel_start + 1];
                let b = image[pixel_start + 2];
                // written as BGR
                self.writer.write_u8(b)?;
                self.writer.write_u8(g)?;
                self.writer.write_u8(r)?;
                // alpha is never written as it's not widely supported
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_rgba(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                let r = image[pixel_start];
                let g = image[pixel_start + 1];
                let b = image[pixel_start + 2];
                let a = image[pixel_start + 3];
                // written as BGRA
                self.writer.write_u8(b)?;
                self.writer.write_u8(g)?;
                self.writer.write_u8(r)?;
                self.writer.write_u8(a)?;
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_gray(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        // write grayscale palette
        for val in 0..256 {
            // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
            let val = val as u8;
            self.writer.write_u8(val)?;
            self.writer.write_u8(val)?;
            self.writer.write_u8(val)?;
            self.writer.write_u8(0)?;
        }

        // write image data
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                // color value is equal to the palette index
                self.writer.write_u8(image[pixel_start])?;
                // alpha is never written as it's not widely supported
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
        for _ in 0..row_pad_size {
            self.writer.write_u8(0)?;
        }

        Ok(())
    }
}

impl<'a, W: Write> ImageEncoder for BMPEncoder<'a, W> {
    fn write_image(
        mut self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: color::ColorType,
    ) -> ImageResult<()> {
        self.encode(buf, width, height, color_type)
    }
}

fn get_unsupported_error_message(c: color::ColorType) -> String {
    format!(
        "Unsupported color type {:?}.  Supported types: RGB(8), RGBA(8), Gray(8), GrayA(8).",
        c
    )
}

/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
fn get_pixel_info(c: color::ColorType) -> io::Result<(u32, u32, u32)> {
    let sizes = match c {
        color::ColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, 0),
        color::ColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, 0),
        color::ColorType::L8 => (BITMAPINFOHEADER_SIZE, 1, 256),
        color::ColorType::La8 => (BITMAPINFOHEADER_SIZE, 1, 256),
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                &get_unsupported_error_message(c)[..],
            ))
        }
    };

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BmpDecoder;
    use super::BMPEncoder;
    use color::ColorType;
    use image::ImageDecoder;
    use std::io::Cursor;

    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ColorType) -> Vec<u8> {
        let mut encoded_data = Vec::new();
        {
            let mut encoder = BMPEncoder::new(&mut encoded_data);
            encoder
                .encode(&image, width, height, c)
                .expect("could not encode image");
        }

        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");

        let mut buf = vec![0; decoder.total_bytes() as usize];
        decoder.read_image(&mut buf).expect("failed to decode");
        buf
    }

    #[test]
    fn round_trip_single_pixel_rgb() {
        let image = [255u8, 0, 0]; // single red pixel
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgb8);
        assert_eq!(3, decoded.len());
        assert_eq!(255, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
    }

    #[test]
    fn huge_files_return_error() {
        let mut encoded_data = Vec::new();
        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
        let mut encoder = BMPEncoder::new(&mut encoded_data);
        let result = encoder.encode(&image, 40_000, 40_000, ColorType::Rgb8);
        assert!(result.is_err());
    }

    #[test]
    fn round_trip_single_pixel_rgba() {
        let image = [1, 2, 3, 4];
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgba8);
        assert_eq!(&decoded[..], &image[..]);
    }

    #[test]
    fn round_trip_3px_rgb() {
        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
        let _decoded = round_trip_image(&image, 3, 3, ColorType::Rgb8);
    }

    #[test]
    fn round_trip_gray() {
        let image = [0u8, 1, 2]; // 3 pixels
        let decoded = round_trip_image(&image, 3, 1, ColorType::L8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }

    #[test]
    fn round_trip_graya() {
        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
        let decoded = round_trip_image(&image, 1, 3, ColorType::La8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }
}
<commit_msg>Refactored using ok_or. Checked also for the last add.<commit_after>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::{self, Write};

use color;
use image::{ImageEncoder, ImageError, ImageResult};

const BITMAPFILEHEADER_SIZE: u32 = 14;
const BITMAPINFOHEADER_SIZE: u32 = 40;
const BITMAPV4HEADER_SIZE: u32 = 108;

/// The representation of a BMP encoder.
pub struct BMPEncoder<'a, W: 'a> {
    writer: &'a mut W,
}

impl<'a, W: Write + 'a> BMPEncoder<'a, W> {
    /// Create a new encoder that writes its output to ```w```.
    pub fn new(w: &'a mut W) -> Self {
        BMPEncoder { writer: w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```.
    pub fn encode(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
    ) -> ImageResult<()> {
        let bmp_header_size = BITMAPFILEHEADER_SIZE;

        let (dib_header_size, written_pixel_size, palette_color_count) = get_pixel_info(c)?;
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes
        let image_size = width
            .checked_mul(height)
            .ok_or(ImageError::DimensionError)?
            .checked_mul(written_pixel_size)
            .ok_or(ImageError::DimensionError)?
            .checked_add(height * row_pad_size)
            .ok_or(ImageError::DimensionError)?;

        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

        // write BMP header
        self.writer.write_u8(b'B')?;
        self.writer.write_u8(b'M')?;
        self.writer.write_u32::<LittleEndian>(file_size)?; // file size
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 1
        self.writer.write_u16::<LittleEndian>(0)?; // reserved 2
        self.writer
            .write_u32::<LittleEndian>(bmp_header_size + dib_header_size + palette_size)?; // image data offset

        // write DIB header
        self.writer.write_u32::<LittleEndian>(dib_header_size)?;
        self.writer.write_i32::<LittleEndian>(width as i32)?;
        self.writer.write_i32::<LittleEndian>(height as i32)?;
        self.writer.write_u16::<LittleEndian>(1)?; // color planes
        self.writer
            .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; // bits per pixel
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(3)?; // compression method - bitfields
        } else {
            self.writer.write_u32::<LittleEndian>(0)?; // compression method - no compression
        }
        self.writer.write_u32::<LittleEndian>(image_size)?;
        self.writer.write_i32::<LittleEndian>(0)?; // horizontal ppm
        self.writer.write_i32::<LittleEndian>(0)?; // vertical ppm
        self.writer.write_u32::<LittleEndian>(palette_color_count)?;
        self.writer.write_u32::<LittleEndian>(0)?; // all colors are important
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            self.writer.write_u32::<LittleEndian>(0xff << 16)?; // red mask
            self.writer.write_u32::<LittleEndian>(0xff << 8)?; // green mask
            self.writer.write_u32::<LittleEndian>(0xff << 0)?; // blue mask
            self.writer.write_u32::<LittleEndian>(0xff << 24)?; // alpha mask
            self.writer.write_u32::<LittleEndian>(0x73524742)?; // colorspace - sRGB

            // endpoints (3x3) and gamma (3)
            for _ in 0..12 {
                self.writer.write_u32::<LittleEndian>(0)?;
            }
        }

        // write image data
        match c {
            color::ColorType::Rgb8 => {
                self.encode_rgb(image, width, height, row_pad_size, 3)?
            }
            color::ColorType::Rgba8 => {
                self.encode_rgba(image, width, height, row_pad_size, 4)?
            }
            color::ColorType::L8 => {
                self.encode_gray(image, width, height, row_pad_size, 1)?
            }
            color::ColorType::La8 => {
                self.encode_gray(image, width, height, row_pad_size, 2)?
            }
            _ => {
                return Err(ImageError::IoError(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                )))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                let r = image[pixel_start];
                let g = image[pixel_start + 1];
                let b = image[pixel_start + 2];
                // written as BGR
                self.writer.write_u8(b)?;
                self.writer.write_u8(g)?;
                self.writer.write_u8(r)?;
                // alpha is never written as it's not widely supported
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_rgba(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                let r = image[pixel_start];
                let g = image[pixel_start + 1];
                let b = image[pixel_start + 2];
                let a = image[pixel_start + 3];
                // written as BGRA
                self.writer.write_u8(b)?;
                self.writer.write_u8(g)?;
                self.writer.write_u8(r)?;
                self.writer.write_u8(a)?;
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn encode_gray(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        // write grayscale palette
        for val in 0..256 {
            // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
            let val = val as u8;
            self.writer.write_u8(val)?;
            self.writer.write_u8(val)?;
            self.writer.write_u8(val)?;
            self.writer.write_u8(0)?;
        }

        // write image data
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                // color value is equal to the palette index
                self.writer.write_u8(image[pixel_start])?;
                // alpha is never written as it's not widely supported
            }

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
        for _ in 0..row_pad_size {
            self.writer.write_u8(0)?;
        }

        Ok(())
    }
}

impl<'a, W: Write> ImageEncoder for BMPEncoder<'a, W> {
    fn write_image(
        mut self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: color::ColorType,
    ) -> ImageResult<()> {
        self.encode(buf, width, height, color_type)
    }
}

fn get_unsupported_error_message(c: color::ColorType) -> String {
    format!(
        "Unsupported color type {:?}.  Supported types: RGB(8), RGBA(8), Gray(8), GrayA(8).",
        c
    )
}

/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
fn get_pixel_info(c: color::ColorType) -> io::Result<(u32, u32, u32)> {
    let sizes = match c {
        color::ColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, 0),
        color::ColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, 0),
        color::ColorType::L8 => (BITMAPINFOHEADER_SIZE, 1, 256),
        color::ColorType::La8 => (BITMAPINFOHEADER_SIZE, 1, 256),
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                &get_unsupported_error_message(c)[..],
            ))
        }
    };

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BmpDecoder;
    use super::BMPEncoder;
    use color::ColorType;
    use image::ImageDecoder;
    use std::io::Cursor;

    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ColorType) -> Vec<u8> {
        let mut encoded_data = Vec::new();
        {
            let mut encoder = BMPEncoder::new(&mut encoded_data);
            encoder
                .encode(&image, width, height, c)
                .expect("could not encode image");
        }

        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");

        let mut buf = vec![0; decoder.total_bytes() as usize];
        decoder.read_image(&mut buf).expect("failed to decode");
        buf
    }

    #[test]
    fn round_trip_single_pixel_rgb() {
        let image = [255u8, 0, 0]; // single red pixel
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgb8);
        assert_eq!(3, decoded.len());
        assert_eq!(255, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
    }

    #[test]
    fn huge_files_return_error() {
        let mut encoded_data = Vec::new();
        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
        let mut encoder = BMPEncoder::new(&mut encoded_data);
        let result = encoder.encode(&image, 40_000, 40_000, ColorType::Rgb8);
        assert!(result.is_err());
    }

    #[test]
    fn round_trip_single_pixel_rgba() {
        let image = [1, 2, 3, 4];
        let decoded = round_trip_image(&image, 1, 1, ColorType::Rgba8);
        assert_eq!(&decoded[..], &image[..]);
    }

    #[test]
    fn round_trip_3px_rgb() {
        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
        let _decoded = round_trip_image(&image, 3, 3, ColorType::Rgb8);
    }

    #[test]
    fn round_trip_gray() {
        let image = [0u8, 1, 2]; // 3 pixels
        let decoded = round_trip_image(&image, 3, 1, ColorType::L8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }

    #[test]
    fn round_trip_graya() {
        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
        let decoded = round_trip_image(&image, 1, 3, ColorType::La8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }
}
<|endoftext|>


<commit_before>use num::bigint::{ToBigUint, RandBigInt, BigUint};
use num::{Zero, One};
use num::integer::Integer;
use rand::thread_rng;

/// Cryptographically useful extensions to the provided BigUint functionality.
pub trait BigUintCrypto {
    /// Find the next prime from the current BigUint
    fn next_prime(&self) -> BigUint;

    /// use the extended euclidean algorithm to solve for (g,x,y) given (a,b) such that
    /// g = gcd(a,b) = a*x + b*y.
    fn gcdext(&self, other: &BigUint) -> (BigUint, BigUint, BigUint);

    /// Is this number a prime number. Uses a probablistic function to determine primality.
    fn is_prime(n: &BigUint) -> bool;

    /// perform the function (base^exponent) % modulus using exponentiation by sqauring
    fn mod_exp(base: &BigUint, exponent: &BigUint, modulus: &BigUint) -> BigUint;
}

impl BigUintCrypto for BigUint {
    fn next_prime(&self) -> BigUint {
        let one: BigUint = One::one();
        let two = 2.to_biguint().unwrap();
        let mut next_prime = self.clone();
        if &next_prime % &two == Zero::zero() {
            next_prime = &next_prime + &one;
        } else {
            next_prime = &next_prime + &two;
        }
        while !BigUint::is_prime(&next_prime) {
            next_prime = &next_prime + &two;
        }
        next_prime
    }

    fn gcdext(&self, other: &BigUint) -> (BigUint, BigUint, BigUint) {

        (Zero::zero(), Zero::zero(), Zero::zero())
    }

    fn is_prime(n: &BigUint) -> bool {
        let two = 2.to_biguint().unwrap();
        let three = 3.to_biguint().unwrap();
        if *n == three || *n == two {
            return true;
        }
        if *n < two || n % two == Zero::zero() {
            return false;
        }
        miller_rabin(n, 100)
    }

    fn mod_exp(base: &BigUint, exponent: &BigUint, modulus: &BigUint) -> BigUint {
        let zero = Zero::zero();
        let one: BigUint = One::one();
        let two = &one + &one;
        let mut result: BigUint = One::one();
        let mut base_acc = base.clone();
        let mut exp_acc = exponent.clone();
        while exp_acc > zero {
            if (&exp_acc % &two) == one {
                result = (result * &base_acc) % modulus;
            }
            exp_acc = exp_acc >> 1;
            base_acc = (&base_acc * &base_acc) % modulus;
        }
        result
    }
}

/// n must be greater than 3 and k indicates the number of rounds
fn miller_rabin(n: &BigUint, k: usize) -> bool{
    let one: BigUint = One::one();
    let two: BigUint = &one + &one;
    let mut d: BigUint = n - &One::one();
    let mut s: BigUint = Zero::zero();
    while d.is_even() {
        d = d >> 1;
        s = s + &One::one();
    }

    for _ in 0..k {
        //println!("loop {} of {}", j, k);
        let a = thread_rng().gen_biguint_range(&two, &(n - &two));
        let mut x = mod_exp(&a, &d, n);
        //let mut x = two.clone();
        if (x == one) || (x == (n - &one)) {
            continue;
        }

        // Use a while loop instead of for here because range does not accept BigUint
        let mut i: BigUint = Zero::zero();
        loop  {
            x = mod_exp(&x, &two, n);
            if x == one || i == (&s - &one) {
                return false;
            }
            if x == (n - &one) {
                break;
            }
            i = i + &one;
        }
    }
    true
}

fn mod_exp(base: &BigUint, exponent: &BigUint, modulus: &BigUint) -> BigUint {
    let zero = Zero::zero();
    let one: BigUint = One::one();
    let two = &one + &one;
    let mut result: BigUint = One::one();
    let mut base_acc = base.clone();
    let mut exp_acc = exponent.clone();
    while exp_acc > zero {
        if (&exp_acc % &two) == one {
            result = (result * &base_acc) % modulus;
        }
        exp_acc = exp_acc >> 1;
        base_acc = (&base_acc * &base_acc) % modulus;
    }
    result
}

#[cfg(test)]
mod test_BigUint_crypto {
    use super::{BigUintCrypto, mod_exp};
    use num::bigint::{RandBigInt, BigUint};
    use std::num::FromPrimitive;
    use num::One;
    use rand::thread_rng;

    #[test]
    fn next_prime_test() {
        let test_num = BigUint::
        parse_bytes("4829837983753984028472098472089547098728675098723407520875258".as_bytes(), 10).unwrap();

        let expected_next = BigUint::
        parse_bytes("4829837983753984028472098472089547098728675098723407520875297".as_bytes(), 10).unwrap();

        assert!(test_num.next_prime() == expected_next);
    }

    #[test]
    fn mod_exp_test() {
        let base = BigUint::from_isize(4).unwrap();
        let exponent = BigUint::from_isize(13).unwrap();
        let modulus = BigUint::from_isize(497).unwrap();
        let expected_result = BigUint::from_isize(445).unwrap();

        assert!(mod_exp(&base, &exponent, &modulus) == expected_result);
    }

    #[test]
    fn is_prime_test() {
        let known_prime = BigUint::
        parse_bytes("4829837983753984028472098472089547098728675098723407520875297".as_bytes(), 10).unwrap();

        assert!(BigUint::is_prime(&known_prime));
    }

    #[test]
    #[should_fail]
    fn is_prime_test_failuire() {
        let not_prime = BigUint::
        parse_bytes("359709793871987301975987296195681798740165298740176567105918720469720137416098423"
        .as_bytes(), 10).unwrap();

        assert!(BigUint::is_prime(&not_prime));
    }

}
<commit_msg>Make miller_rabin parallel.<commit_after>use num::bigint::{ToBigUint, RandBigInt, BigUint};
use num::{Zero, One};
use num::integer::Integer;
use rand::thread_rng;
use std::sync::{Arc, mpsc};
use std::thread;

/// Cryptographically useful extensions to the provided BigUint functionality.
pub trait BigUintCrypto {
    /// Find the next prime from the current BigUint
    fn next_prime(&self) -> BigUint;

    /// use the extended euclidean algorithm to solve for (g,x,y) given (a,b) such that
    /// g = gcd(a,b) = a*x + b*y.
    fn gcdext(&self, other: &BigUint) -> (BigUint, BigUint, BigUint);

    /// Is this number a prime number. Uses a probablistic function to determine primality.
    fn is_prime(n: &BigUint) -> bool;

    /// perform the function (base^exponent) % modulus using exponentiation by sqauring
    fn mod_exp(base: &BigUint, exponent: &BigUint, modulus: &BigUint) -> BigUint;
}

impl BigUintCrypto for BigUint {
    fn next_prime(&self) -> BigUint {
        let one: BigUint = One::one();
        let two = 2.to_biguint().unwrap();
        let mut next_prime = self.clone();
        if &next_prime % &two == Zero::zero() {
            next_prime = &next_prime + &one;
        } else {
            next_prime = &next_prime + &two;
        }
        while !BigUint::is_prime(&next_prime) {
            next_prime = &next_prime + &two;
        }
        next_prime
    }

    fn gcdext(&self, other: &BigUint) -> (BigUint, BigUint, BigUint) {

        (Zero::zero(), Zero::zero(), Zero::zero())
    }

    fn is_prime(n: &BigUint) -> bool {
        let two = 2.to_biguint().unwrap();
        let three = 3.to_biguint().unwrap();
        if *n == three || *n == two {
            return true;
        }
        if *n < two || n % two == Zero::zero() {
            return false;
        }
        miller_rabin(n, 100)
    }

    fn mod_exp(base: &BigUint, exponent: &BigUint, modulus: &BigUint) -> BigUint {
        let zero = Zero::zero();
        let one: BigUint = One::one();
        let two = &one + &one;
        let mut result: BigUint = One::one();
        let mut base_acc = base.clone();
        let mut exp_acc = exponent.clone();
        while exp_acc > zero {
            if (&exp_acc % &two) == one {
                result = (result * &base_acc) % modulus;
            }
            exp_acc = exp_acc >> 1;
            base_acc = (&base_acc * &base_acc) % modulus;
        }
        result
    }
}

/// n must be greater than 3 and k indicates the number of rounds
fn miller_rabin(n: &BigUint, k: usize) -> bool{
    let one: BigUint = One::one();
    let (tx, rx) = mpsc::channel();

    let mut d: BigUint = n - &One::one();
    let mut s: BigUint = Zero::zero();
    while d.is_even() {
        d = d >> 1;
        s = s + &one;
    }
    let shared_n = Arc::new(n.clone());
    let shared_d = Arc::new(d);
    let shared_s = Arc::new(s);

    // miller rabin lends itself to being concurrent since a is completely random
    // here we spawn multiple threads to help speed up the process
    for _ in 0..8 {
        let tx = tx.clone();
        //let thread_n = n.clone();
        let shared_d = shared_d.clone();
        let shared_s = shared_s.clone();
        let shared_n = shared_n.clone();
        thread::spawn(move || {
            let result = miller_rabin_thread(&shared_n, &shared_d, &shared_s, k/8);
            tx.send(result);
            });
    }

    for _ in 0..8 {
        if !rx.recv().ok().expect("A thread failed") {
            return false;
        }
    }
    true
}

fn miller_rabin_thread(n: &BigUint, d: &BigUint, s: &BigUint, k: usize) -> bool {
    let one: BigUint = One::one();
    let two: BigUint = &one + &one;

    for _ in 0..k {
        //println!("loop {} of {}", j, k);
        let a = thread_rng().gen_biguint_range(&two, &(n - &two));
        let mut x = mod_exp(&a, d, n);
        //let mut x = two.clone();
        if (x == one) || (x == (n - &one)) {
            continue;
        }

        // Use a while loop instead of for here because range does not accept BigUint
        let mut i: BigUint = Zero::zero();
        loop  {
            x = mod_exp(&x, &two, n);
            if x == one || i == (s - &one) {
                return false;
            }
            if x == (n - &one) {
                break;
            }
            i = i + &one;
        }
    }
    true
}

fn mod_exp(base: &BigUint, exponent: &BigUint, modulus: &BigUint) -> BigUint {
    let zero = Zero::zero();
    let one: BigUint = One::one();
    let two = &one + &one;
    let mut result: BigUint = One::one();
    let mut base_acc = base.clone();
    let mut exp_acc = exponent.clone();
    while exp_acc > zero {
        if (&exp_acc % &two) == one {
            result = (result * &base_acc) % modulus;
        }
        exp_acc = exp_acc >> 1;
        base_acc = (&base_acc * &base_acc) % modulus;
    }
    result
}

#[cfg(test)]
mod test_BigUint_crypto {
    use super::{BigUintCrypto, mod_exp};
    use num::bigint::{RandBigInt, BigUint};
    use std::num::FromPrimitive;
    use num::One;
    use rand::thread_rng;

    #[test]
    fn next_prime_test() {
        let test_num = BigUint::
        parse_bytes("4829837983753984028472098472089547098728675098723407520875258".as_bytes(), 10).unwrap();

        let expected_next = BigUint::
        parse_bytes("4829837983753984028472098472089547098728675098723407520875297".as_bytes(), 10).unwrap();

        assert!(test_num.next_prime() == expected_next);
    }

    #[test]
    fn mod_exp_test() {
        let base = BigUint::from_isize(4).unwrap();
        let exponent = BigUint::from_isize(13).unwrap();
        let modulus = BigUint::from_isize(497).unwrap();
        let expected_result = BigUint::from_isize(445).unwrap();

        assert!(mod_exp(&base, &exponent, &modulus) == expected_result);
    }

    #[test]
    fn is_prime_test() {
        let known_prime = BigUint::
        parse_bytes("4829837983753984028472098472089547098728675098723407520875297".as_bytes(), 10).unwrap();

        assert!(BigUint::is_prime(&known_prime));
    }

    #[test]
    #[should_fail]
    fn is_prime_test_failuire() {
        let not_prime = BigUint::
        parse_bytes("359709793871987301975987296195681798740165298740176567105918720469720137416098423"
        .as_bytes(), 10).unwrap();

        assert!(BigUint::is_prime(&not_prime));
    }

}
<|endoftext|>


<commit_before>use bindings::*;
use std::slice::bytes::copy_memory;

static KEY_BYTES: uint = 32;
static NONCE_BYTES: uint = 24;
static ZERO_BYTES: uint = 32;

pub struct SecretKey ([u8, ..KEY_BYTES]);

impl SecretKey {
    pub fn from_str(str: &str) -> SecretKey {
        SecretKey::from_slice(str.as_bytes())
    }

    pub fn from_slice(slice: &[u8]) -> SecretKey {
        assert!(slice.len() <= KEY_BYTES);

        let mut sized = [0u8, ..KEY_BYTES];
        copy_memory(sized, slice);

        SecretKey(sized)
    }
}

pub struct SecretMsg {
    pub nonce: [u8, ..NONCE_BYTES],
    pub cipher: Vec<u8>
}

pub struct SecretBox {
    sk: SecretKey
}

impl SecretBox {
    pub fn new(sk: SecretKey) -> SecretBox {
        SecretBox { sk: sk }
    }

    pub fn encrypt(&self, msg: &[u8]) -> SecretMsg {
        unsafe {
            let mut stretched: Vec<u8> = Vec::from_elem(ZERO_BYTES, 0u8);
            stretched.push_all(msg);

            let mut nonce = [0u8, ..NONCE_BYTES];
            randombytes(nonce.as_mut_ptr(), NONCE_BYTES as u64);

            let SecretKey(sk) = self.sk;

            let mut cipher = Vec::from_elem(stretched.len(), 0u8);

            match crypto_secretbox(cipher.as_mut_ptr(),
                                   stretched.as_ptr(),
                                   stretched.len() as u64,
                                   nonce.as_ptr(),
                                   sk.as_ptr()) {
                0 => SecretMsg { nonce: nonce, cipher: cipher },
                _ => fail!("crypto_secretbox failed")
            }
        }
    }

    pub fn decrypt(&self, msg: &SecretMsg) -> Vec<u8> {
        let &SecretMsg { ref nonce, ref cipher } = msg;

        unsafe {
            let mut msg = Vec::from_elem(cipher.len(), 0u8);
            let SecretKey(sk) = self.sk;

            match crypto_secretbox_open(msg.as_mut_ptr(),
                                        cipher.as_ptr(),
                                        cipher.len() as u64,
                                        nonce.as_ptr(),
                                        sk.as_ptr()) {
                0 => Vec::from_slice(msg.slice(ZERO_BYTES, msg.len())),
                _ => fail!("crypto_secretbox_open failed")
            }
        }
    }
}

#[test]
fn test_secretbox_sanity() {
    for i in range(0 as uint, 16) {
        let msg = Vec::from_elem(i * 4, i as u8);

        let sb = SecretBox::new(SecretKey::from_str("some secret key"));
        let SecretMsg { nonce, cipher } = sb.encrypt(msg.as_slice());

        println!("enc:\t{}\nnonce:\t{}", cipher, Vec::from_slice(nonce));

        let decr = sb.decrypt(&SecretMsg { nonce: nonce, cipher: cipher });

        println!("dec:\t{}", decr);

        assert!(decr == msg);
    }
}

#[test]
fn test_secretbox_uniqueness() {
    let msg = Vec::from_elem(128, 0x53u8);

    let box1 = SecretBox::new(SecretKey::from_str("1"));
    let box2 = SecretBox::new(SecretKey::from_str(""));

    let SecretMsg { nonce: n1, cipher: c1 } = box1.encrypt(msg.as_slice());
    let SecretMsg { nonce: n2, cipher: c2 } = box2.encrypt(msg.as_slice());

    assert!(n1 != n2);
    assert!(c1 != c2);
    assert!(c1 != msg);
    assert!(c2 != msg);
}
<commit_msg>Add some docs for secretbox<commit_after>//! The `secretbox` module symmetrically encrypts given plaintext and
//! then uses a one time authenticator to ensure tamper-resistance.
//!
//! In other words, this uses an encrypt-then-MAC scheme.
//!
//! TODO: document me.

use bindings::*;
use std::slice::bytes::copy_memory;

/// Size of shared secret key used for symmetric encryption.
pub static KEY_BYTES: uint = 32;
/// Size of the nonce value.
pub static NONCE_BYTES: uint = 24;
/// Size of the zero padding applied to each message.
pub static ZERO_BYTES: uint = 32;

/// Shared secret key. Must be `<= KEY_BYTES` bytes long.
pub struct SecretKey ([u8, ..KEY_BYTES]);

impl SecretKey {
    /// Generate a secret key from the bytes of a given string.
    pub fn from_str(str: &str) -> SecretKey {
        SecretKey::from_slice(str.as_bytes())
    }

    /// Generate a secret key from a slice (turn a slice into a sized slice).
    pub fn from_slice(slice: &[u8]) -> SecretKey {
        assert!(slice.len() <= KEY_BYTES);

        let mut sized = [0u8, ..KEY_BYTES];
        copy_memory(sized, slice);

        SecretKey(sized)
    }
}

/// Encapsulates both the nonce value and cipher text returned by `encrypt`.
pub struct SecretMsg {
    pub nonce: [u8, ..NONCE_BYTES],
    pub cipher: Vec<u8>
}

/// TODO: Document me
pub struct SecretBox {
    sk: SecretKey
}

impl SecretBox {
    /// Create a new SecretBox with the provided key
    pub fn new(sk: SecretKey) -> SecretBox {
        SecretBox { sk: sk }
    }

    /// Using this box's secret key, symmetrically encrypt the given message.
    ///
    /// A random nonce value will be securely generated and returned
    /// as part of the response.
    pub fn encrypt(&self, msg: &[u8]) -> SecretMsg {
        unsafe {
            let mut stretched: Vec<u8> = Vec::from_elem(ZERO_BYTES, 0u8);
            stretched.push_all(msg);

            let mut nonce = [0u8, ..NONCE_BYTES];
            randombytes(nonce.as_mut_ptr(), NONCE_BYTES as u64);

            let SecretKey(sk) = self.sk;

            let mut cipher = Vec::from_elem(stretched.len(), 0u8);

            // TODO: Better error handling
            match crypto_secretbox(cipher.as_mut_ptr(),
                                   stretched.as_ptr(),
                                   stretched.len() as u64,
                                   nonce.as_ptr(),
                                   sk.as_ptr()) {
                0 => SecretMsg { nonce: nonce, cipher: cipher },
                _ => fail!("crypto_secretbox failed")
            }
        }
    }

    /// Using this box's secret key, decrypt the given ciphertext into
    /// plain text.
    pub fn decrypt(&self, msg: &SecretMsg) -> Vec<u8> {
        let &SecretMsg { ref nonce, ref cipher } = msg;

        unsafe {
            let mut msg = Vec::from_elem(cipher.len(), 0u8);
            let SecretKey(sk) = self.sk;

            // TODO: Error handling, this can fail in non-fatal ways
            //       (MAC validation fails etc.)
            match crypto_secretbox_open(msg.as_mut_ptr(),
                                        cipher.as_ptr(),
                                        cipher.len() as u64,
                                        nonce.as_ptr(),
                                        sk.as_ptr()) {
                0 => Vec::from_slice(msg.slice(ZERO_BYTES, msg.len())),
                _ => fail!("crypto_secretbox_open failed")
            }
        }
    }
}

#[test]
fn test_secretbox_sanity() {
    for i in range(0 as uint, 16) {
        let msg = Vec::from_elem(i * 4, i as u8);

        let sb = SecretBox::new(SecretKey::from_str("some secret key"));
        let SecretMsg { nonce, mut cipher } = sb.encrypt(msg.as_slice());

        println!("enc:\t{}\nnonce:\t{}", cipher, Vec::from_slice(nonce));

        let decr = sb.decrypt(&SecretMsg { nonce: nonce, cipher: cipher });

        println!("dec:\t{}", decr);

        assert!(decr == msg);
    }
}

#[test]
fn test_secretbox_uniqueness() {
    let msg = Vec::from_elem(128, 0x53u8);

    let box1 = SecretBox::new(SecretKey::from_str("1"));
    let box2 = SecretBox::new(SecretKey::from_str(""));

    let SecretMsg { nonce: n1, cipher: c1 } = box1.encrypt(msg.as_slice());
    let SecretMsg { nonce: n2, cipher: c2 } = box2.encrypt(msg.as_slice());

    assert!(n1 != n2);
    assert!(c1 != c2);
    assert!(c1 != msg);
    assert!(c2 != msg);
}
<|endoftext|>