<commit_before>use clap::Clap;
use std::cmp::Ordering;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)$").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace("(", "")
        .replace(")", "")
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_")
        .replace("#", "");
    if !col.chars().next().map(char::is_alphabetic).unwrap_or(true) {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            match row.len().cmp(&normalized_cols.len()) {
                Ordering::Less => {
                    for _ in 0..normalized_cols.len() - row.len() {
                        row.push_field("");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(&row)).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({}) from {:?}",
        num_rows,
        table_name,
        normalized_cols.join(", "),
        path,
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else if line.starts_with(".schema") {
        _handle_query(
            conn,
            "SELECT sql AS schema FROM sqlite_master WHERE name like 't%'",
        )
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

fn install_udfs(c: &mut rusqlite::Connection) -> Result<(), Box<dyn std::error::Error>> {
    c.create_scalar_function(
        "regexp_extract",
        3,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            let replacement = ctx.get::<Box<str>>(2)?;

            let caps = match re.captures(&value) {
                Some(caps) => caps,
                None => return Ok("".to_string()),
            };
            let mut dest = String::new();
            caps.expand(&replacement, &mut dest);
            Ok(dest)
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    type Hint = String;

    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(Clap)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,
    #[clap(long, about = "Use ';' as the delimiter for the CSV")]
    semicolon: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    let delim = match (opts.comma, opts.pipe, opts.tab, opts.semicolon) {
        (true, false, false, false) | (false, false, false, false) => b',',
        (false, true, false, false) => b'|',
        (false, false, true, false) => b'\t',
        (false, false, false, true) => b';',
        _ => {
            eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
            std::process::exit(1);
        }
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        "length",
        "coalesce",
        "regexp_extract",
        ".export",
        ".schema",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    install_udfs(&mut conn)?;

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("", ""),
            ("abc", "abc"),
            ("(S)AO", "sao"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
            ("COMBO#", "combo"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<commit_msg>remove unnecesary borrow (#571)<commit_after>use clap::Clap;
use std::cmp::Ordering;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)$").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace("(", "")
        .replace(")", "")
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_")
        .replace("#", "");
    if !col.chars().next().map(char::is_alphabetic).unwrap_or(true) {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            match row.len().cmp(&normalized_cols.len()) {
                Ordering::Less => {
                    for _ in 0..normalized_cols.len() - row.len() {
                        row.push_field("");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(&row)).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({}) from {:?}",
        num_rows,
        table_name,
        normalized_cols.join(", "),
        path,
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else if line.starts_with(".schema") {
        _handle_query(
            conn,
            "SELECT sql AS schema FROM sqlite_master WHERE name like 't%'",
        )
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

fn install_udfs(c: &mut rusqlite::Connection) -> Result<(), Box<dyn std::error::Error>> {
    c.create_scalar_function(
        "regexp_extract",
        3,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            let replacement = ctx.get::<Box<str>>(2)?;

            let caps = match re.captures(&value) {
                Some(caps) => caps,
                None => return Ok("".to_string()),
            };
            let mut dest = String::new();
            caps.expand(&replacement, &mut dest);
            Ok(dest)
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    type Hint = String;

    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(Clap)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,
    #[clap(long, about = "Use ';' as the delimiter for the CSV")]
    semicolon: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    let delim = match (opts.comma, opts.pipe, opts.tab, opts.semicolon) {
        (true, false, false, false) | (false, false, false, false) => b',',
        (false, true, false, false) => b'|',
        (false, false, true, false) => b'\t',
        (false, false, false, true) => b';',
        _ => {
            eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
            std::process::exit(1);
        }
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        "length",
        "coalesce",
        "regexp_extract",
        ".export",
        ".schema",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    install_udfs(&mut conn)?;

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("", ""),
            ("abc", "abc"),
            ("(S)AO", "sao"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
            ("COMBO#", "combo"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<|endoftext|>


<commit_before>// rust version of breakfast tool
// BreakFast is a toolkit for detecting chromosomal rearrangements
// based on RNA-seq data.

use std::env;
use std::mem;
use std::thread;
use std::sync::Arc;
use std::path::Path;
use std::io::prelude::*;
use std::fs;
use std::os::unix::io::AsRawFd;
use std::os::unix::io::FromRawFd;
use std::fs::File;
use std::process::{Command, Stdio};
use std::collections::HashMap;
use std::io::{BufReader, BufRead, Read, Write};
use rust_htslib::bam;
use rust_htslib::bam::Read as HTSRead;

extern crate clap;
extern crate bio;
extern crate rust_htslib;
extern crate subprocess;
extern crate regex;

#[macro_use]
extern crate shells;

use subprocess::{Exec, Popen, PopenConfig, Redirection};
use regex::Regex;

mod cli;
mod lib;

// struct for having breakfast options
struct BfOptions {
    anchor_len :    usize,
    max_frag_len:   i32,
    min_mapq:       i32,
    orientation:    &'static str,
    all_reads:      bool,
    discard_duplicates: &'static str,
    min_reads:      &'static str,
    freq_above:     i32,
}

struct Evidence {
	chr: String,
	pos: u32,
	strand: bool,
	mchr: String,
	mpos: u32,
	mstrand: bool,
	sequence: String,
	frag_id: String,
	signature: String     // Breakpoint signature
}

fn main() {
  let defaults = BfOptions{ anchor_len : 0,
                            max_frag_len: 5000,
                            min_mapq: 15,
                            orientation:  "fr",
                            all_reads:  false,
                            discard_duplicates: "both-ends",
                            min_reads:    "0-0-0",
                            freq_above: 0};

  let mut anchor_len: usize   = defaults.anchor_len;
  let mut max_frag_len: i32 = defaults.max_frag_len;
  let mut min_mapq: i32     = defaults.min_mapq;
  let mut orientation       = String::from(defaults.orientation);
  let mut all_reads: bool   = defaults.all_reads;
  let mut discard_duplicates = String::from(defaults.discard_duplicates);
  let mut min_reads     = String::from(defaults.min_reads);
  let mut freq_above: i32 = defaults.freq_above;

  //arguments from src/cli.rs
  let matches = cli::build_cli().get_matches();
  //arguments from detect subcommand
  if let ("detect", Some(detect)) = matches.subcommand() {
      let bam_file = String::from(detect.value_of("bam_file").unwrap());
      let genome   = String::from(detect.value_of("genome").unwrap());

      if detect.is_present("anchor-len") {
          anchor_len = detect.value_of("anchor-len").unwrap().parse().unwrap();
      }

      if detect.is_present("orientation") {
          orientation = String::from(detect.value_of("orientation").unwrap());
      }
      detect_discordant_reads(bam_file, genome, anchor_len);
  }
}

// DETECT //
//======= //

fn detect_discordant_pairs(sam_path: String, out_prefix: String, max_frag_len: i32,
	min_mapq: i32, orientation: String) {
  println!("{:?}\t{:?}\t{:?}\t{:?}\t{:?}", sam_path, out_prefix, max_frag_len, min_mapq, orientation);

  let mut N = 0;
  let tmp = out_prefix.clone();
  let given_path = Path::new(&tmp);
  let mut sort_tmp_dir = Path::new("./");
  if given_path.is_absolute()  {
        sort_tmp_dir = given_path.parent().unwrap();
  }

  let extn = String::from(".discordant_pairs.tsv.gz");
  let out  = out_prefix + &extn;
  let mut file = File::create("test.txt").unwrap();

  lib::mkdir(&tmp);
  println!("Searching for discordant read pairs ...");

  let (code, stdout, stderr) = sh!("sam discordant pairs --min-mapq={} {} {} | sort -k1,1 -T {:?}", min_mapq, sam_path, max_frag_len, sort_tmp_dir);

  let mut prev : Vec<&str> = Vec::new();
  for line in stdout.lines() {
      let mut cols: Vec<&str> = line.split('\t').collect();
      if cols.len() < 9 { continue; }

      let r1 = Regex::new(r"N").unwrap();
      let r2 = Regex::new(r"S").unwrap();

      if r1.is_match(cols[5]) || r2.is_match(cols[5]) { continue; }
      else if cols[0].ends_with("/1") || cols[0].ends_with("/2") {
          let tmp : Vec<&str> = cols[0].split('/').collect();
          cols[0] = tmp[0];
      }
      if prev.len() <= 0 { prev = cols.clone(); }
      else if prev.len() > 0 && cols[0] != prev[0] { prev = cols.clone(); }

      let flags = &cols[1].parse::<i32>().unwrap();
      let mut chr   = String::from(cols[2]);
      let mut mchr  = String::from(prev[2]);
      let mut strand = '-';
      let mut mstrand = '-';
      let mut pos  = cols[3].parse::<i32>().unwrap();
      let mut mpos = prev[3].parse::<i32>().unwrap();
      let mut rlen = cols[9].clone().len() as i32;
      let mut mrlen= prev[9].clone().len() as i32;

      if chr.starts_with("chr") {  chr = chr; }
      else { let prefix = String::from("chr"); chr = prefix + &chr;}

      if mchr.starts_with("chr") {  chr = chr ;}
      else { let prefix = String::from("chr"); mchr = prefix + &mchr;}

      if chr.contains("chrM") || mchr.contains("chrM") { continue;}

      if orientation == "fr" {
        if chr > mchr || (chr == mchr && pos > mpos) {
          mem::swap(&mut chr, &mut mchr);
          mem::swap(&mut pos, &mut mpos);
          mem::swap(&mut rlen, &mut mrlen);
          mem::swap(&mut strand, &mut mstrand);
        }
        if mstrand == '+' { mstrand = '-';}
        else { mstrand = '+'; }
      }

      else if orientation == "rf" {
        if chr > mchr || (chr == mchr && pos > mpos) {
          mem::swap(&mut chr, &mut mchr);
          mem::swap(&mut pos, &mut mpos);
          mem::swap(&mut rlen, &mut mrlen);
          mem::swap(&mut strand, &mut mstrand);
        }
        if strand == '+' { strand = '-';}
        else { strand = '+'; }
      }

      else if orientation == "ff" {
        if chr > mchr || (chr == mchr && pos > mpos) {
          mem::swap(&mut chr, &mut mchr);
          mem::swap(&mut pos, &mut mpos);
          mem::swap(&mut rlen, &mut mrlen);

          if mstrand == '-' { mstrand = '+';} else { mstrand = '-'; }
          if strand  == '-' { strand  = '+';} else { strand  = '-'; }
          mem::swap(&mut strand, &mut mstrand);
        }
      } else { panic!("Unsupported read orientation detected."); }

      if strand  == '-' { pos  +=  rlen - 1;}
      if mstrand == '-' { mpos += mrlen - 1;}

      let mut out = String::new();
      out.push_str(&chr);out.push_str(&"\t");
      out.push_str(&strand.to_string());out.push_str(&"\t");
      out.push_str(&pos.to_string());out.push_str(&"\t");
      out.push_str(&mchr);out.push_str(&"\t");
      out.push_str(&mstrand.to_string());out.push_str(&"\t");
      out.push_str(&mpos.to_string());out.push_str(&"\t-\n");
      file.write_all(out.as_bytes());
      N += 1;
    }
    let (wcode, wstdout, wstderr) = sh!("gzip -c {:?}.txt > {}", given_path, out);
    let mut rm_tmp = tmp.to_string() + ".txt";
    fs::remove_file(rm_tmp);
    println!("Found {:?} discordant mate pairs.", N);
}

fn detect_discordant_reads(sam_path: String, genome_path: String, anchor_len: usize) {

	let fastq = bio::io::fasta::Reader::from_file(format!("{}.fa", genome_path)).unwrap();
	println!("Reading reference genome into memory...");
	let mut genome = HashMap::new();
	for entry in fastq.records() {
		let chr = entry.unwrap();
		genome.insert(chr.id().unwrap().to_owned(), chr.seq().to_owned());
	}

    println!("Splitting unaligned reads into {} bp anchors and aligning against the genome...", anchor_len);

    let mut bowtie = Command::new("bowtie")
		.args(&["-f", "-p1", "-v0", "-m1", "-B1", "--suppress", "5,6,7,8", &genome_path, "-"])
        .stdin(Stdio::piped()).stdout(Stdio::piped())
        .spawn().unwrap_or_else(|e| {
            panic!("Bowtie failed running!");
        });

	let mut bowtie_in = unsafe { File::from_raw_fd(bowtie.stdin.unwrap().as_raw_fd()) };
	let mut bowtie_out = unsafe { BufReader::new(File::from_raw_fd(bowtie.stdout.unwrap().as_raw_fd())) };

	thread::spawn(move || {
		let bam = bam::Reader::from_path(&sam_path).unwrap();
		let mut R = 0;
		for r in bam.records() {
			let read = r.unwrap();
			if read.is_unmapped() == false { continue; }
			if read.seq().len() < anchor_len * 2 { continue; }
			// TODO: Extract anchors from both ends of read and write in
			// interleaved FASTA format to the stdin of Bowtie.
	        R += 1;
	        let seq = String::from_utf8(read.seq().as_bytes()).unwrap();
	        //let tail = seq.len() - anchor_len;
	        write!(bowtie_in, ">{}_1\n{}", R, &seq[..anchor_len]);
            println!(">{:?}_1\n{}", R, &seq);
	    }
    });

    //println!("{}", child_out.recv().unwrap());
    //TODO: bowtie is not recognizing fata format
    // check the formatting
    // write and feed a tmpfile to bowtie maybe!?
    //write!(bowtie_stdin, "{:?}", child_out.recv().unwrap());

     let mut bowtie_results = String::new();
       match bowtie_out.read_to_string(&mut bowtie_results) {
           Err(why) => panic!("bowtie not running!"),
           Ok(_) => print!("bowtie worked!"),
       }

    println!(" There are {} lines\t", bowtie_results);

    // match fas.stdin.unwrap().write_all(samres.as_bytes()) {
    //     Err(why) => panic!("samtools results not reaching fasta"),
    //     Ok(_) => print!("fasta recvied input from samtools"),
    // }

    // let mut res = String::new();
    //   match fas.stdout.unwrap().read_to_string(&mut res) {
    //       Err(why) => panic!("fasta not running!"),
    //       Ok(_) => print!("fasta subcommand worked!"),
    //   }
    //   println!("{:?}", res.len());
}
<commit_msg>dry run<commit_after>// rust version of breakfast tool
// BreakFast is a toolkit for detecting chromosomal rearrangements
// based on RNA-seq data.

use std::env;
use std::mem;
use std::thread;
use std::sync::Arc;
use std::path::Path;
use std::io::prelude::*;
use std::fs;
use std::os::unix::io::AsRawFd;
use std::os::unix::io::FromRawFd;
use std::fs::File;
use std::process::{Command, Stdio};
use std::collections::HashMap;
use std::io::{BufReader, BufRead, Read, Write};
use rust_htslib::bam;
use rust_htslib::bam::Read as HTSRead;

extern crate clap;
extern crate bio;
extern crate rust_htslib;
extern crate subprocess;
extern crate regex;

#[macro_use]
extern crate shells;

use subprocess::{Exec, Popen, PopenConfig, Redirection};
use regex::Regex;

mod cli;
mod lib;

// struct for having breakfast options
struct BfOptions {
    anchor_len :    usize,
    max_frag_len:   i32,
    min_mapq:       i32,
    orientation:    &'static str,
    all_reads:      bool,
    discard_duplicates: &'static str,
    min_reads:      &'static str,
    freq_above:     i32,
}

struct Evidence {
	chr: String,
	pos: u32,
	strand: bool,
	mchr: String,
	mpos: u32,
	mstrand: bool,
	sequence: String,
	frag_id: String,
	signature: String     // Breakpoint signature
}

fn main() {
  let defaults = BfOptions{ anchor_len : 0,
                            max_frag_len: 5000,
                            min_mapq: 15,
                            orientation:  "fr",
                            all_reads:  false,
                            discard_duplicates: "both-ends",
                            min_reads:    "0-0-0",
                            freq_above: 0};

  let mut anchor_len: usize   = defaults.anchor_len;
  let mut max_frag_len: i32 = defaults.max_frag_len;
  let mut min_mapq: i32     = defaults.min_mapq;
  let mut orientation       = String::from(defaults.orientation);
  let mut all_reads: bool   = defaults.all_reads;
  let mut discard_duplicates = String::from(defaults.discard_duplicates);
  let mut min_reads     = String::from(defaults.min_reads);
  let mut freq_above: i32 = defaults.freq_above;

  //arguments from src/cli.rs
  let matches = cli::build_cli().get_matches();
  //arguments from detect subcommand
  if let ("detect", Some(detect)) = matches.subcommand() {
      let bam_file = String::from(detect.value_of("bam_file").unwrap());
      let genome   = String::from(detect.value_of("genome").unwrap());

      if detect.is_present("anchor-len") {
          anchor_len = detect.value_of("anchor-len").unwrap().parse().unwrap();
      }

      if detect.is_present("orientation") {
          orientation = String::from(detect.value_of("orientation").unwrap());
      }
      detect_discordant_reads(bam_file, genome, anchor_len);
  }
}

// DETECT //
//======= //

fn detect_discordant_pairs(sam_path: String, out_prefix: String, max_frag_len: i32,
	min_mapq: i32, orientation: String) {
  println!("{:?}\t{:?}\t{:?}\t{:?}\t{:?}", sam_path, out_prefix, max_frag_len, min_mapq, orientation);

  let mut N = 0;
  let tmp = out_prefix.clone();
  let given_path = Path::new(&tmp);
  let mut sort_tmp_dir = Path::new("./");
  if given_path.is_absolute()  {
        sort_tmp_dir = given_path.parent().unwrap();
  }

  let extn = String::from(".discordant_pairs.tsv.gz");
  let out  = out_prefix + &extn;
  let mut file = File::create("test.txt").unwrap();

  lib::mkdir(&tmp);
  println!("Searching for discordant read pairs ...");

  let (code, stdout, stderr) = sh!("sam discordant pairs --min-mapq={} {} {} | sort -k1,1 -T {:?}", min_mapq, sam_path, max_frag_len, sort_tmp_dir);

  let mut prev : Vec<&str> = Vec::new();
  for line in stdout.lines() {
      let mut cols: Vec<&str> = line.split('\t').collect();
      if cols.len() < 9 { continue; }

      let r1 = Regex::new(r"N").unwrap();
      let r2 = Regex::new(r"S").unwrap();

      if r1.is_match(cols[5]) || r2.is_match(cols[5]) { continue; }
      else if cols[0].ends_with("/1") || cols[0].ends_with("/2") {
          let tmp : Vec<&str> = cols[0].split('/').collect();
          cols[0] = tmp[0];
      }
      if prev.len() <= 0 { prev = cols.clone(); }
      else if prev.len() > 0 && cols[0] != prev[0] { prev = cols.clone(); }

      let flags = &cols[1].parse::<i32>().unwrap();
      let mut chr   = String::from(cols[2]);
      let mut mchr  = String::from(prev[2]);
      let mut strand = '-';
      let mut mstrand = '-';
      let mut pos  = cols[3].parse::<i32>().unwrap();
      let mut mpos = prev[3].parse::<i32>().unwrap();
      let mut rlen = cols[9].clone().len() as i32;
      let mut mrlen= prev[9].clone().len() as i32;

      if chr.starts_with("chr") {  chr = chr; }
      else { let prefix = String::from("chr"); chr = prefix + &chr;}

      if mchr.starts_with("chr") {  chr = chr ;}
      else { let prefix = String::from("chr"); mchr = prefix + &mchr;}

      if chr.contains("chrM") || mchr.contains("chrM") { continue;}

      if orientation == "fr" {
        if chr > mchr || (chr == mchr && pos > mpos) {
          mem::swap(&mut chr, &mut mchr);
          mem::swap(&mut pos, &mut mpos);
          mem::swap(&mut rlen, &mut mrlen);
          mem::swap(&mut strand, &mut mstrand);
        }
        if mstrand == '+' { mstrand = '-';}
        else { mstrand = '+'; }
      }

      else if orientation == "rf" {
        if chr > mchr || (chr == mchr && pos > mpos) {
          mem::swap(&mut chr, &mut mchr);
          mem::swap(&mut pos, &mut mpos);
          mem::swap(&mut rlen, &mut mrlen);
          mem::swap(&mut strand, &mut mstrand);
        }
        if strand == '+' { strand = '-';}
        else { strand = '+'; }
      }

      else if orientation == "ff" {
        if chr > mchr || (chr == mchr && pos > mpos) {
          mem::swap(&mut chr, &mut mchr);
          mem::swap(&mut pos, &mut mpos);
          mem::swap(&mut rlen, &mut mrlen);

          if mstrand == '-' { mstrand = '+';} else { mstrand = '-'; }
          if strand  == '-' { strand  = '+';} else { strand  = '-'; }
          mem::swap(&mut strand, &mut mstrand);
        }
      } else { panic!("Unsupported read orientation detected."); }

      if strand  == '-' { pos  +=  rlen - 1;}
      if mstrand == '-' { mpos += mrlen - 1;}

      let mut out = String::new();
      out.push_str(&chr);out.push_str(&"\t");
      out.push_str(&strand.to_string());out.push_str(&"\t");
      out.push_str(&pos.to_string());out.push_str(&"\t");
      out.push_str(&mchr);out.push_str(&"\t");
      out.push_str(&mstrand.to_string());out.push_str(&"\t");
      out.push_str(&mpos.to_string());out.push_str(&"\t-\n");
      file.write_all(out.as_bytes());
      N += 1;
    }
    let (wcode, wstdout, wstderr) = sh!("gzip -c {:?}.txt > {}", given_path, out);
    let mut rm_tmp = tmp.to_string() + ".txt";
    fs::remove_file(rm_tmp);
    println!("Found {:?} discordant mate pairs.", N);
}

fn detect_discordant_reads(sam_path: String, genome_path: String, anchor_len: usize) {

	let fastq = bio::io::fasta::Reader::from_file(format!("{}.fa", genome_path)).unwrap();
	println!("Reading reference genome into memory...");
	let mut genome = HashMap::new();
	for entry in fastq.records() {
		let chr = entry.unwrap();
		genome.insert(chr.id().unwrap().to_owned(), chr.seq().to_owned());
	}

    println!("Splitting unaligned reads into {} bp anchors and aligning against the genome...", anchor_len);
    let mut bowtie = Command::new("bowtie")
		.args(&["-f", "-p1", "-v0", "-m1", "-B1", "--suppress", "5,6,7,8", &genome_path, "-"])
        .stdin(Stdio::piped()).stdout(Stdio::piped())
        .spawn().unwrap_or_else(|e| {
            panic!("Bowtie failed running!");
        });

	let mut bowtie_in  = unsafe { File::from_raw_fd(bowtie.stdin.unwrap().as_raw_fd()) };
	let mut bowtie_out = unsafe { BufReader::new(File::from_raw_fd(bowtie.stdout.unwrap().as_raw_fd())) };

	thread::spawn(move || {
		let bam = bam::Reader::from_path(&sam_path).unwrap();
		let mut R = 0;
		for r in bam.records() {
			let read = r.unwrap();
			if read.is_unmapped() == false { continue; }
			if read.seq().len() < anchor_len * 2 { continue; }
			// TODO: Extract anchors from both ends of read and write in
			// interleaved FASTA format to the stdin of Bowtie.
	        R += 1;
	        let seq = String::from_utf8(read.seq().as_bytes()).unwrap();
	        //let tail = seq.len() - anchor_len;
	        write!(bowtie_in, ">{}_1\n{}", R, &seq[..anchor_len].to_string());
         //   println!(">{:?}_1\n{}", R, &seq[..anchor_len].to_string());
	    }
    });

    //println!("{}", child_out.recv().unwrap());
    //TODO: bowtie is not recognizing fata format
    // check the formatting
    // write and feed a tmpfile to bowtie maybe!?
    //write!(bowtie_stdin, "{:?}", child_out.recv().unwrap());

    let mut bowtie_results = String::new();
       match bowtie_out.read_to_string(&mut bowtie_results) {
           Err(why) => panic!("bowtie not running!"),
           Ok(_) => print!("bowtie worked!"),
       }

    println!(" There are {} lines\t", bowtie_results);

    // match fas.stdin.unwrap().write_all(samres.as_bytes()) {
    //     Err(why) => panic!("samtools results not reaching fasta"),
    //     Ok(_) => print!("fasta recvied input from samtools"),
    // }

    // let mut res = String::new();
    //   match fas.stdout.unwrap().read_to_string(&mut res) {
    //       Err(why) => panic!("fasta not running!"),
    //       Ok(_) => print!("fasta subcommand worked!"),
    //   }
    //   println!("{:?}", res.len());
}
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate crc;

use std::fs::File;
use std::io::{Read, Result, Error, Seek, SeekFrom};
use std::io::ErrorKind::Other;

use byteorder::{LittleEndian, ByteOrder};

use crc::{crc32, Hasher32};

const LABEL_SCAN_SECTORS: usize = 4;
const MYPATH: &'static str = "/dev/mapper/luks-291e9250-200e-4c36-8e5e-66aa257669ed";
const ID_LEN: usize = 32;
const MDA_MAGIC: &'static str = "\x20\x4c\x56\x4d\x32\x20\x78\x5b\x35\x41\x25\x72\x30\x4e\x2a\x3e";
const INITIAL_CRC: u32 = 0xf597a6cf;
const SECTOR_SIZE: usize = 512;

#[derive(Debug)]
struct Label {
    id: String,
    sector: u64,
    crc: u32,
    offset: u32,
    label: String,
}

#[derive(Debug)]
struct PvArea {
    offset: u64,
    size: u64,
}

#[derive(Debug)]
struct PvHeader {
    uuid: String,
    size: u64, // in bytes
    ext_version: u32,
    ext_flags: u32,
    data_areas: Vec<PvArea>,
    metadata_areas: Vec<PvArea>,
    bootloader_areas: Vec<PvArea>,
}

fn get_label(buf: &[u8]) -> Result<Label> {
    for x in 0..LABEL_SCAN_SECTORS {
        let sec_buf = &buf[x*SECTOR_SIZE..x*SECTOR_SIZE+SECTOR_SIZE];
        if &sec_buf[..8] == b"LABELONE" {
            let crc = LittleEndian::read_u32(&sec_buf[16..20]);
            crc32_ok(crc, &sec_buf[20..SECTOR_SIZE]);

            let sector = LittleEndian::read_u64(&sec_buf[8..16]);
            if sector != x as u64 {
                println!("sector field is {} in sector {}", sector, x);
            }

            return Ok(Label{
                id: String::from_utf8_lossy(&sec_buf[..8]).into_owned(),
                sector: sector,
                crc: crc,
                offset: LittleEndian::read_u32(&sec_buf[20..24]) + (x*SECTOR_SIZE as usize) as u32,
                label: String::from_utf8_lossy(&sec_buf[24..32]).into_owned(),
            })
        }
    }

    Err(Error::new(Other, "bad"))
}

struct PvAreaIter<'a> {
    area: &'a[u8],
}

fn iter_pv_area<'a>(buf: &'a[u8]) -> PvAreaIter<'a> {
    PvAreaIter { area: buf }
}

impl<'a> Iterator for PvAreaIter<'a> {
    type Item = PvArea;

    fn next (&mut self) -> Option<PvArea> {
        let off = LittleEndian::read_u64(&self.area[..8]);
        let size = LittleEndian::read_u64(&self.area[8..16]);

        if off == 0 {
            None
        }
        else {
            self.area = &self.area[16..];
            Some(PvArea {
                offset: off,
                size: size,
            })
        }
    }
}

//
// PV HEADER LAYOUT:
// - static header (uuid and size)
// - 0+ data areas
// - blank entry
// - 0+ metadata areas
// - blank entry
// - 8 bytes of pvextension header
// - if version > 0
//   - list of bootloader areas
//
fn get_pv_header(buf: &[u8]) -> Result<PvHeader> {

    let mut da_buf = &buf[ID_LEN+8..];

    let da_vec: Vec<_> = iter_pv_area(da_buf).collect();

    // move slice past any actual entries plus blank
    // terminating entry
    da_buf = &da_buf[(da_vec.len()+1)*16..];

    let md_vec: Vec<_> = iter_pv_area(da_buf).collect();

    da_buf = &da_buf[(md_vec.len()+1)*16..];

    let ext_version = LittleEndian::read_u32(&da_buf[..4]);
    let mut ext_flags = 0;
    let mut ba_vec = Vec::new();

    if ext_version != 0 {
        ext_flags = LittleEndian::read_u32(&da_buf[4..8]);

        da_buf = &da_buf[8..];

        ba_vec = iter_pv_area(da_buf).collect();
    }

    Ok(PvHeader{
        uuid: String::from_utf8_lossy(&buf[..ID_LEN]).into_owned(),
        size: LittleEndian::read_u64(&buf[ID_LEN..ID_LEN+8]),
        ext_version: ext_version,
        ext_flags: ext_flags,
        data_areas: da_vec,
        metadata_areas: md_vec,
        bootloader_areas: ba_vec,
    })
}

fn crc32_ok(val: u32, buf: &[u8]) -> bool {
    let mut digest = crc32::Digest::new(INITIAL_CRC);
    digest.write(&buf);
    let crc32 = digest.sum32();
    if val != crc32 {
        println!("CRC32: input {:x} != calculated {:x}", val, crc32);
    }
    val == crc32
}

fn parse_mda_header(buf: &[u8]) -> () {

    let crc1 = LittleEndian::read_u32(&buf[..4]);

    // TODO: why is this failing?
    crc32_ok(crc1, &buf[4..512]);
}

fn find_stuff(path: &str) -> Result<Label> {

    let mut f = try!(File::open(path));

    let mut buf = vec![0; LABEL_SCAN_SECTORS * 512];

    try!(f.read(&mut buf));

    let label = try!(get_label(&buf));

    let pvheader = try!(get_pv_header(&buf[label.offset as usize..]));

    println!("pvheader {:?}", pvheader);

    for md in &pvheader.metadata_areas {
        try!(f.seek(SeekFrom::Start(md.offset)));

        println!("AA {} {}", md.offset, md.size);

        let mut buf = vec![0; md.size as usize];

        try!(f.read(&mut buf));

        parse_mda_header(&buf);
    }

    return Ok(label);
}

fn main() {

    let label = find_stuff(MYPATH).unwrap();

    println!("{:?}", label);
}
<commit_msg>check MDA magic<commit_after>extern crate byteorder;
extern crate crc;

use std::fs::File;
use std::io::{Read, Result, Error, Seek, SeekFrom};
use std::io::ErrorKind::Other;

use byteorder::{LittleEndian, ByteOrder};

use crc::{crc32, Hasher32};

const LABEL_SCAN_SECTORS: usize = 4;
const MYPATH: &'static str = "/dev/mapper/luks-291e9250-200e-4c36-8e5e-66aa257669ed";
const ID_LEN: usize = 32;
const MDA_MAGIC: &'static [u8] = b"\x20\x4c\x56\x4d\x32\x20\x78\x5b\x35\x41\x25\x72\x30\x4e\x2a\x3e";
const INITIAL_CRC: u32 = 0xf597a6cf;
const SECTOR_SIZE: usize = 512;

#[derive(Debug)]
struct Label {
    id: String,
    sector: u64,
    crc: u32,
    offset: u32,
    label: String,
}

#[derive(Debug)]
struct PvArea {
    offset: u64,
    size: u64,
}

#[derive(Debug)]
struct PvHeader {
    uuid: String,
    size: u64, // in bytes
    ext_version: u32,
    ext_flags: u32,
    data_areas: Vec<PvArea>,
    metadata_areas: Vec<PvArea>,
    bootloader_areas: Vec<PvArea>,
}

fn get_label(buf: &[u8]) -> Result<Label> {
    for x in 0..LABEL_SCAN_SECTORS {
        let sec_buf = &buf[x*SECTOR_SIZE..x*SECTOR_SIZE+SECTOR_SIZE];
        if &sec_buf[..8] == b"LABELONE" {
            let crc = LittleEndian::read_u32(&sec_buf[16..20]);
            crc32_ok(crc, &sec_buf[20..SECTOR_SIZE]);

            let sector = LittleEndian::read_u64(&sec_buf[8..16]);
            if sector != x as u64 {
                println!("sector field is {} in sector {}", sector, x);
            }

            return Ok(Label{
                id: String::from_utf8_lossy(&sec_buf[..8]).into_owned(),
                sector: sector,
                crc: crc,
                offset: LittleEndian::read_u32(&sec_buf[20..24]) + (x*SECTOR_SIZE as usize) as u32,
                label: String::from_utf8_lossy(&sec_buf[24..32]).into_owned(),
            })
        }
    }

    Err(Error::new(Other, "bad"))
}

struct PvAreaIter<'a> {
    area: &'a[u8],
}

fn iter_pv_area<'a>(buf: &'a[u8]) -> PvAreaIter<'a> {
    PvAreaIter { area: buf }
}

impl<'a> Iterator for PvAreaIter<'a> {
    type Item = PvArea;

    fn next (&mut self) -> Option<PvArea> {
        let off = LittleEndian::read_u64(&self.area[..8]);
        let size = LittleEndian::read_u64(&self.area[8..16]);

        if off == 0 {
            None
        }
        else {
            self.area = &self.area[16..];
            Some(PvArea {
                offset: off,
                size: size,
            })
        }
    }
}

//
// PV HEADER LAYOUT:
// - static header (uuid and size)
// - 0+ data areas
// - blank entry
// - 0+ metadata areas
// - blank entry
// - 8 bytes of pvextension header
// - if version > 0
//   - list of bootloader areas
//
fn get_pv_header(buf: &[u8]) -> Result<PvHeader> {

    let mut da_buf = &buf[ID_LEN+8..];

    let da_vec: Vec<_> = iter_pv_area(da_buf).collect();

    // move slice past any actual entries plus blank
    // terminating entry
    da_buf = &da_buf[(da_vec.len()+1)*16..];

    let md_vec: Vec<_> = iter_pv_area(da_buf).collect();

    da_buf = &da_buf[(md_vec.len()+1)*16..];

    let ext_version = LittleEndian::read_u32(&da_buf[..4]);
    let mut ext_flags = 0;
    let mut ba_vec = Vec::new();

    if ext_version != 0 {
        ext_flags = LittleEndian::read_u32(&da_buf[4..8]);

        da_buf = &da_buf[8..];

        ba_vec = iter_pv_area(da_buf).collect();
    }

    Ok(PvHeader{
        uuid: String::from_utf8_lossy(&buf[..ID_LEN]).into_owned(),
        size: LittleEndian::read_u64(&buf[ID_LEN..ID_LEN+8]),
        ext_version: ext_version,
        ext_flags: ext_flags,
        data_areas: da_vec,
        metadata_areas: md_vec,
        bootloader_areas: ba_vec,
    })
}

fn crc32_ok(val: u32, buf: &[u8]) -> bool {
    let mut digest = crc32::Digest::new(INITIAL_CRC);
    digest.write(&buf);
    let crc32 = digest.sum32();
    if val != crc32 {
        println!("CRC32: input {:x} != calculated {:x}", val, crc32);
    }
    val == crc32
}

fn parse_mda_header(buf: &[u8]) -> () {

    let crc1 = LittleEndian::read_u32(&buf[..4]);

    // TODO: why is this failing?
    crc32_ok(crc1, &buf[4..512]);

    if &buf[4..20] != MDA_MAGIC {
        println!("'{}' doesn't match '{}'", String::from_utf8_lossy(&buf[4..20]),
                 String::from_utf8_lossy(MDA_MAGIC));
    }
}

fn find_stuff(path: &str) -> Result<Label> {

    let mut f = try!(File::open(path));

    let mut buf = vec![0; LABEL_SCAN_SECTORS * 512];

    try!(f.read(&mut buf));

    let label = try!(get_label(&buf));

    let pvheader = try!(get_pv_header(&buf[label.offset as usize..]));

    for md in &pvheader.metadata_areas {
        try!(f.seek(SeekFrom::Start(md.offset)));

        let mut buf = vec![0; md.size as usize];

        try!(f.read(&mut buf));

        parse_mda_header(&buf);
    }

    return Ok(label);
}

fn main() {

    let label = find_stuff(MYPATH).unwrap();

    println!("{}", label.id);
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015 Jonathan Anderson
 *
 * Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
 * http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
 * http://opensource.org/licenses/MIT>, at your option. This file may not be
 * copied, modified, or distributed except according to those terms.
 */

extern crate docopt;
extern crate rshark;
extern crate rustc_serialize;
extern crate pcap;

use docopt::Docopt;
use rshark::Protocol;

// TODO: use docopt_macros once rust-lang/rust#28089 is resolved
const USAGE: &'static str = "
Usage: rshark [options] <source>
       rshark (--help | --version)

Options:
    -f, --filter         BFP filter (see http://biot.com/capstats/bpf.html)
    -h, --help           Show this message
    -p, --promiscuous    Listen to all packets
    -s, --snaplen=<len>  Bytes to capture from each packet [default: 5000]
    -t, --timeout=<ms>   Packet read timeout, in ms [default: 10]
    -v, --version        Show the version of rshark
";

const VERSION: Option<&'static str> = option_env!("CARGO_PKG_VERSION");

#[derive(RustcDecodable)]
struct Args {
    arg_source: String,
    flag_filter: String,
    flag_snaplen: i32,
    flag_timeout: i32,
    flag_promiscuous: bool,
    flag_version: bool,
}

type PcapResult = Result<pcap::Capture<pcap::Activated>, pcap::Error>;

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.argv(std::env::args()).decode())
        .unwrap_or_else(|e| e.exit())
        ;

    if args.flag_version {
        println!["rshark v{}", VERSION.unwrap_or("<unknown>")];
        return;
    }

    let protocol = rshark::ethernet::Ethernet;

    let result = open_capture(&args)
        .map(|mut c| {
            let mut count = 0;

            while let Ok(packet) = c.next() {
                println!("\n\n{}-B packet:", packet.data.len());

                match protocol.dissect(packet.data) {
                    Ok(dissected) => print!["{}", dissected.pretty_print(1)],
                    Err(e) => println!["Error: {}", e],
                }

                count += 1;
            }

            count
        })
        ;

    match result {
        Ok(packet_count) => println!["Processed {} packets", packet_count],
        Err(e) => {
            println!["{}", e];
            std::process::exit(1);
        },
    }
}

fn open_capture(args: &Args) -> PcapResult {
    let device = try![open_device(args)];

    let capture = match device {
        Some(d) => Ok(d),
        None => open_file(&args.arg_source),
    };

    capture.and_then(|mut c| {
        try![c.filter(&args.flag_filter)];
        Ok(c)
    })
}

fn open_device(args: &Args)
        -> Result<Option<pcap::Capture<pcap::Activated>>, pcap::Error> {

    match pcap::Device::list() {
        Ok(devices) => {
            for d in devices {
                if d.name == args.arg_source {
                    return pcap::Capture::from_device(d)
                        .map(|d| d.promisc(args.flag_promiscuous)
                                  .rfmon(args.flag_promiscuous)
                                  .snaplen(args.flag_snaplen)
                                  .timeout(args.flag_timeout))
                        .and_then(|d| d.open())
                        .map(|c| c.into())
                        .map(Some)
                        ;
                }
            };

            Ok(None)
        },
        Err(e) => Err(e),
    }
}

fn open_file(filename: &str) -> PcapResult {
    std::fs::metadata(filename)
        .map_err(|e| pcap::Error::PcapError(format!["{}", e]))
        .and_then(|f|
            if f.is_file() {
                pcap::Capture::from_file(filename)
                    .map(|c| c.into())
            } else {
                Err(pcap::Error::PcapError(
                        format!["{} is not a file or interface", filename]))
            }
        )
}
<commit_msg>Make version printing unconditional.<commit_after>/*
 * Copyright 2015 Jonathan Anderson
 *
 * Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
 * http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
 * http://opensource.org/licenses/MIT>, at your option. This file may not be
 * copied, modified, or distributed except according to those terms.
 */

extern crate docopt;
extern crate rshark;
extern crate rustc_serialize;
extern crate pcap;

use docopt::Docopt;
use rshark::Protocol;

// TODO: use docopt_macros once rust-lang/rust#28089 is resolved
const USAGE: &'static str = "
Usage: rshark [options] <source>
       rshark (--help | --version)

Options:
    -f, --filter         BFP filter (see http://biot.com/capstats/bpf.html)
    -h, --help           Show this message
    -p, --promiscuous    Listen to all packets
    -s, --snaplen=<len>  Bytes to capture from each packet [default: 5000]
    -t, --timeout=<ms>   Packet read timeout, in ms [default: 10]
    -v, --version        Show the version of rshark
";

const VERSION: Option<&'static str> = option_env!("CARGO_PKG_VERSION");

#[derive(RustcDecodable)]
struct Args {
    arg_source: String,
    flag_filter: String,
    flag_snaplen: i32,
    flag_timeout: i32,
    flag_promiscuous: bool,
    flag_version: bool,
}

type PcapResult = Result<pcap::Capture<pcap::Activated>, pcap::Error>;

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.argv(std::env::args()).decode())
        .unwrap_or_else(|e| e.exit())
        ;

    println!["rshark v{}", VERSION.unwrap_or("<unknown>")];

    if args.flag_version {
        return;
    }

    let protocol = rshark::ethernet::Ethernet;

    let result = open_capture(&args)
        .map(|mut c| {
            let mut count = 0;

            while let Ok(packet) = c.next() {
                println!("\n\n{}-B packet:", packet.data.len());

                match protocol.dissect(packet.data) {
                    Ok(dissected) => print!["{}", dissected.pretty_print(1)],
                    Err(e) => println!["Error: {}", e],
                }

                count += 1;
            }

            count
        })
        ;

    match result {
        Ok(packet_count) => println!["Processed {} packets", packet_count],
        Err(e) => {
            println!["{}", e];
            std::process::exit(1);
        },
    }
}

fn open_capture(args: &Args) -> PcapResult {
    let device = try![open_device(args)];

    let capture = match device {
        Some(d) => Ok(d),
        None => open_file(&args.arg_source),
    };

    capture.and_then(|mut c| {
        try![c.filter(&args.flag_filter)];
        Ok(c)
    })
}

fn open_device(args: &Args)
        -> Result<Option<pcap::Capture<pcap::Activated>>, pcap::Error> {

    match pcap::Device::list() {
        Ok(devices) => {
            for d in devices {
                if d.name == args.arg_source {
                    return pcap::Capture::from_device(d)
                        .map(|d| d.promisc(args.flag_promiscuous)
                                  .rfmon(args.flag_promiscuous)
                                  .snaplen(args.flag_snaplen)
                                  .timeout(args.flag_timeout))
                        .and_then(|d| d.open())
                        .map(|c| c.into())
                        .map(Some)
                        ;
                }
            };

            Ok(None)
        },
        Err(e) => Err(e),
    }
}

fn open_file(filename: &str) -> PcapResult {
    std::fs::metadata(filename)
        .map_err(|e| pcap::Error::PcapError(format!["{}", e]))
        .and_then(|f|
            if f.is_file() {
                pcap::Capture::from_file(filename)
                    .map(|c| c.into())
            } else {
                Err(pcap::Error::PcapError(
                        format!["{} is not a file or interface", filename]))
            }
        )
}
<|endoftext|>


<commit_before>extern crate docopt;
extern crate handlebars;
extern crate pulldown_cmark;
extern crate serde;

use docopt::Docopt;
use handlebars::Handlebars;
use serde::Deserialize;
use std::fs;
use std::path::{Path, PathBuf};
use crate::util::RelativeFrom;

mod file_type;
mod util;

// Docopt usage string
static USAGE: &'static str = "
Usage: rust-notes [options] <source> <dest>

Options:
    -b, --base-url BASE     Base URL for site. Should not include hostname.
";

#[derive(Debug, Deserialize)]
struct Args {
    arg_source: String,
    arg_dest: String,
    flag_base_url: Option<String>,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());
    //Generator(args).start();
    match Generator::new(args) {
        Ok(generator) => {
            generator.begin();
        }
        Err(message) => panic!(message),
    }
}

fn cp_dir(source: &Path, dest: &Path) {
    fs::create_dir(dest)
        .ok()
        .expect("Problem copying directory");
    for item in util::walk_dir(source)
        .ok()
        .expect("Problem copying directory")
    {
        let item = item.ok().expect("Problem copying directory").path();
        let item_metadata = fs::metadata(&item)
            .ok()
            .expect("Error fetching file metadata");
        let relative = item.my_relative_from(source).unwrap();
        let dest_path = dest.clone().join(&relative);
        if item_metadata.is_file() {
            fs::copy(&item, &dest_path)
                .ok()
                .expect("Problem copying directory");
        } else {
            fs::create_dir(&dest_path)
                .ok()
                .expect("Problem copying directory");
        }
    }
}

pub struct AppContext<'a> {
    root_source: PathBuf,
    root_dest: PathBuf,
    root_notes: PathBuf,
    handlebars: Handlebars<'a>,
    base_url: String,
}

struct Generator<'a> {
    context: AppContext<'a>,
    file_type_manager: file_type::FileTypeManager,
}

impl<'a> Generator<'a> {
    fn convert(&self, path: &Path) {
        self.file_type_manager
            .create_file_type(path)
            .convert(&self.context);
    }

    pub fn new(args: Args) -> Result<Generator<'a>, &'static str> {
        let source_path = Path::new(&args.arg_source);
        let source_path_metadata = fs::metadata(source_path).ok().expect(&format!(
            "Error fetching file metadata. Source path: {:?}",
            source_path
        ));
        let dest_path = Path::new(&args.arg_dest);
        let dest_path_metadata = fs::metadata(dest_path).ok().expect(&format!(
            "Error fetching file metadata. Dest path: {:?}",
            dest_path
        ));

        if !source_path_metadata.is_dir() {
            return Err("Invalid source path");
        }

        if !dest_path_metadata.is_dir() {
            match fs::create_dir_all(&dest_path) {
                Err(_) => return Err("Cannot create destination directory"),
                _ => (),
            }
        }

        // Validate source
        let notes_source_path = source_path.join("notes");
        let notes_source_path_metadata = fs::metadata(source_path)
            .ok()
            .expect("Error fetching file metadata");

        if !notes_source_path_metadata.is_dir() {
            return Err("Source directory missing required files");
        }

        let base_url = match args.flag_base_url {
            Some(ref base_url) if base_url.is_empty() => None,
            Some(base_url) => {
                let mut result = String::from(base_url.trim_matches('/'));
                result = format!("/{}/", result);
                Some(result)
            }
            None => None,
        };

        let mut context = AppContext {
            root_source: PathBuf::from(source_path),
            root_dest: PathBuf::from(dest_path),
            root_notes: notes_source_path,
            handlebars: Handlebars::new(),
            base_url: base_url.clone().unwrap_or(String::from("/")),
        };

        let file_type_manager = file_type::FileTypeManager::new();
        file_type_manager.initialize_app_context(&mut context)?;

        // Good to go! Let's return something good

        Ok(Generator {
            context: context,
            file_type_manager: file_type_manager,
        })
    }

    pub fn begin(&self) {
        self.clean_dest();
        let assets_source_path = self.context.root_source.join("assets");
        let assets_source_path_metadata = fs::metadata(&assets_source_path)
            .ok()
            .expect("Error fetching file metadata");
        if assets_source_path_metadata.is_dir() {
            let assets_dest_path = self.context.root_dest.join("assets");
            cp_dir(&assets_source_path, &assets_dest_path);
        }
        self.convert(&self.context.root_notes);
        for item in util::walk_dir(&self.context.root_notes).ok().unwrap() {
            self.convert(&item.ok().unwrap().path());
        }
    }

    fn clean_dest(&self) {
        for entry in fs::read_dir(&self.context.root_dest).ok().unwrap() {
            let entry = entry.ok().unwrap();
            let entry_metadata = fs::metadata(entry.path())
                .ok()
                .expect("Error fetching file metadata");
            if entry_metadata.is_file() {
                fs::remove_file(entry.path())
                    .ok()
                    .expect("Could not remove file");
            } else {
                fs::remove_dir_all(entry.path())
                    .ok()
                    .expect("Could not remove directory");
            }
        }
    }
}
<commit_msg>removed extern crate things<commit_after>use crate::util::RelativeFrom;
use docopt::Docopt;
use handlebars::Handlebars;
use serde::Deserialize;
use std::fs;
use std::path::{Path, PathBuf};

mod file_type;
mod util;

// Docopt usage string
static USAGE: &'static str = "
Usage: rust-notes [options] <source> <dest>

Options:
    -b, --base-url BASE     Base URL for site. Should not include hostname.
";

#[derive(Debug, Deserialize)]
struct Args {
    arg_source: String,
    arg_dest: String,
    flag_base_url: Option<String>,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());
    //Generator(args).start();
    match Generator::new(args) {
        Ok(generator) => {
            generator.begin();
        }
        Err(message) => panic!(message),
    }
}

fn cp_dir(source: &Path, dest: &Path) {
    fs::create_dir(dest)
        .ok()
        .expect("Problem copying directory");
    for item in util::walk_dir(source)
        .ok()
        .expect("Problem copying directory")
    {
        let item = item.ok().expect("Problem copying directory").path();
        let item_metadata = fs::metadata(&item)
            .ok()
            .expect("Error fetching file metadata");
        let relative = item.my_relative_from(source).unwrap();
        let dest_path = dest.clone().join(&relative);
        if item_metadata.is_file() {
            fs::copy(&item, &dest_path)
                .ok()
                .expect("Problem copying directory");
        } else {
            fs::create_dir(&dest_path)
                .ok()
                .expect("Problem copying directory");
        }
    }
}

pub struct AppContext<'a> {
    root_source: PathBuf,
    root_dest: PathBuf,
    root_notes: PathBuf,
    handlebars: Handlebars<'a>,
    base_url: String,
}

struct Generator<'a> {
    context: AppContext<'a>,
    file_type_manager: file_type::FileTypeManager,
}

impl<'a> Generator<'a> {
    fn convert(&self, path: &Path) {
        self.file_type_manager
            .create_file_type(path)
            .convert(&self.context);
    }

    pub fn new(args: Args) -> Result<Generator<'a>, &'static str> {
        let source_path = Path::new(&args.arg_source);
        let source_path_metadata = fs::metadata(source_path).ok().expect(&format!(
            "Error fetching file metadata. Source path: {:?}",
            source_path
        ));
        let dest_path = Path::new(&args.arg_dest);
        let dest_path_metadata = fs::metadata(dest_path).ok().expect(&format!(
            "Error fetching file metadata. Dest path: {:?}",
            dest_path
        ));

        if !source_path_metadata.is_dir() {
            return Err("Invalid source path");
        }

        if !dest_path_metadata.is_dir() {
            match fs::create_dir_all(&dest_path) {
                Err(_) => return Err("Cannot create destination directory"),
                _ => (),
            }
        }

        // Validate source
        let notes_source_path = source_path.join("notes");
        let notes_source_path_metadata = fs::metadata(source_path)
            .ok()
            .expect("Error fetching file metadata");

        if !notes_source_path_metadata.is_dir() {
            return Err("Source directory missing required files");
        }

        let base_url = match args.flag_base_url {
            Some(ref base_url) if base_url.is_empty() => None,
            Some(base_url) => {
                let mut result = String::from(base_url.trim_matches('/'));
                result = format!("/{}/", result);
                Some(result)
            }
            None => None,
        };

        let mut context = AppContext {
            root_source: PathBuf::from(source_path),
            root_dest: PathBuf::from(dest_path),
            root_notes: notes_source_path,
            handlebars: Handlebars::new(),
            base_url: base_url.clone().unwrap_or(String::from("/")),
        };

        let file_type_manager = file_type::FileTypeManager::new();
        file_type_manager.initialize_app_context(&mut context)?;

        // Good to go! Let's return something good

        Ok(Generator {
            context: context,
            file_type_manager: file_type_manager,
        })
    }

    pub fn begin(&self) {
        self.clean_dest();
        let assets_source_path = self.context.root_source.join("assets");
        let assets_source_path_metadata = fs::metadata(&assets_source_path)
            .ok()
            .expect("Error fetching file metadata");
        if assets_source_path_metadata.is_dir() {
            let assets_dest_path = self.context.root_dest.join("assets");
            cp_dir(&assets_source_path, &assets_dest_path);
        }
        self.convert(&self.context.root_notes);
        for item in util::walk_dir(&self.context.root_notes).ok().unwrap() {
            self.convert(&item.ok().unwrap().path());
        }
    }

    fn clean_dest(&self) {
        for entry in fs::read_dir(&self.context.root_dest).ok().unwrap() {
            let entry = entry.ok().unwrap();
            let entry_metadata = fs::metadata(entry.path())
                .ok()
                .expect("Error fetching file metadata");
            if entry_metadata.is_file() {
                fs::remove_file(entry.path())
                    .ok()
                    .expect("Could not remove file");
            } else {
                fs::remove_dir_all(entry.path())
                    .ok()
                    .expect("Could not remove directory");
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate num;
extern crate portaudio;
extern crate rtlsdr;

use portaudio as pa;
use rtlsdr::{RTLSDRDevice, RTLSDRError};
use num::complex::Complex;
use std::env;
use std::collections::VecDeque;
use std::f32::consts::PI;
use std::error::Error;

// Threading libraries.
use std::thread;
use std::sync::mpsc::channel;

// Make sure that the buffer size is radix-2, otherwise the read_sync function
// will fail with an error code of -8.
const BUF_SIZE: usize = 131072;

// Most other sample rates fail, but this one works for my particular device.
// I will investigate exactly what's happening here and generate a list of
// possible sample rates (if there are others besides this one).
const SAMPLE_RATE: u32 = 1140000;

// Here is the sample rate of the output waveform we'll try to use.
const SAMPLE_RATE_AUDIO: f64 = 44100.0;

const DEVIATION: u32 = 2500;

const BANDWIDTH: u32 = 200000;

const CHANNELS: i32 = 2;

fn main() {
    match run() {
        Ok(()) => (),
        Err(e) => println!("ERROR: {}", e),
    }
}

struct RTLSDR {
    rtlsdr: RTLSDRDevice,
}

unsafe impl Send for RTLSDR {}

fn run() -> Result<(), Box<Error>> {
    let sdr_index: i32 = env::args().nth(1)
        .expect("Please specify the SDR index")
        .parse::<i32>()?;
    let fm_freq_mhz: f32 = env::args().nth(2)
        .expect("Please specify a center frequency (MHz).")
        .parse::<f32>()?;

    // Initialize the PortAudio class.
    let audio = try!(pa::PortAudio::new());
    let settings =
        try!(audio.default_output_stream_settings(CHANNELS,
                                                  SAMPLE_RATE_AUDIO,
                                                  BUF_SIZE as u32));
    let mut stream = try!(audio.open_blocking_stream(settings));
    try!(stream.start());

    let fm_freq: u32 = (fm_freq_mhz * 1e6) as u32;
    let mut sdr = RTLSDR {
        rtlsdr: init_sdr(sdr_index, fm_freq).unwrap(),
    };

    // TODO: does this gain value make sense based on what I'm doing?
    let gain = (2.0 * PI * DEVIATION as f32 / SAMPLE_RATE_AUDIO as f32).recip();
    let (tx, rx) = channel();

    thread::spawn(move || {
        loop {
            // Read the samples and decimate down to match the sample rate of the
            // audio that'll be going out.
            //
            // TODO: on 88.7 I can make out some instruments, but the audio is still
            // very, very poor. Need to figure out what we need to do here. At a
            // minimum, we need a low-pass filter to filter out the high
            // frequency components that are almost certainly aliasing. In addition,
            // we are underflowing the audio buffer. Not sure if this is a
            // performance issue (release didn't help) or if I need to re-think how
            // I send data to the sound card.
            let tx = tx.clone();
            let bytes = sdr.rtlsdr.read_sync(BUF_SIZE).unwrap();
            let dec_rate = (SAMPLE_RATE as f64/ SAMPLE_RATE_AUDIO) as usize;
            let iq_vec = decimate(read_samples(bytes).unwrap().as_slice(), dec_rate);

            // After decimation, demodulate the signal.
            //
            // TODO: from what I've read there should be some decimation here as
            // well but with narrowband FM, we decimate the signal into
            // nothingness. Need to figure out what we do between wideband and
            // narrowband FM.
            let mut demod_iq = demod_fm(iq_vec, gain);
            tx.send(demod_iq).unwrap();
        }
    });

    loop {
        // Get the write stream and write our samples to the stream.
        let mut demod_iq = rx.recv().unwrap();
        let out_frames = match stream.write_available() {
            Ok(available) => {
                match available {
                    pa::StreamAvailable::Frames(frames) => frames as u32,
                    _ => return Err(Box::new(pa::error::Error::NullCallback)),
                }
            }
            Err(e) => return Err(Box::new(e)),
        };

        let buffer_frames = (demod_iq.len() / CHANNELS as usize) as
                            u32;
        let write_frames = if buffer_frames >= out_frames {
            out_frames
        } else {
            buffer_frames
        };
        let n_write_samples = write_frames as usize * CHANNELS as usize;

        stream.write(write_frames,
                   |output| for ix in 0..n_write_samples as usize {
                       output[ix] = demod_iq.pop_front().unwrap();
                   })?;
    }
}

fn init_sdr(sdr_index: i32, fm_freq: u32) -> Result<RTLSDRDevice, RTLSDRError> {
    let mut sdr = try!(rtlsdr::open(sdr_index));
    try!(sdr.set_center_freq(fm_freq));
    try!(sdr.set_sample_rate(SAMPLE_RATE));
    try!(sdr.set_tuner_bandwidth(BANDWIDTH));
    try!(sdr.set_agc_mode(true));
    try!(sdr.reset_buffer());
    Ok(sdr)
}

fn decimate_queue<T: Copy>(signal: VecDeque<T>, dec_rate: usize) -> VecDeque<T> {
    let mut ix = 0;
    let mut signal_dec = VecDeque::<T>::new();
    while ix < signal.len() {
        signal_dec.push_back(*(signal.get(ix).unwrap()));
        ix += dec_rate;
    }
    signal_dec
}

fn decimate(signal: &[Complex<f32>], dec_rate: usize) -> Vec<Complex<f32>> {
    let mut ix = 0;
    let new_size = (signal.len() / dec_rate + 1) as usize;
    let mut signal_dec = Vec::<Complex<f32>>::with_capacity(new_size);
    while ix < signal.len() {
        signal_dec.push(signal[ix]);
        ix += dec_rate;
    }
    signal_dec
}

fn read_samples(bytes: Vec<u8>) -> Option<Vec<Complex<f32>>> {
    // The docs tell us that RTL-SDR sends us data in bytes alternating between
    // I and Q. I'm not normalizing here, but I might have to in the future.

    // First, check that we've been given an even number of bytes. Not sure
    // what to do if we don't get I and Q.
    let bytes_len = bytes.len();
    if bytes_len % 2 != 0 {
        println!("Received an odd number of samples.");
        return None;

    }
    let mut iq_vec: Vec<Complex<f32>> = Vec::with_capacity(bytes_len / 2);

    // Write the values to the complex value and normalize from [0, 255] to
    // [-1, 1].
    for iq in bytes.chunks(2) {
        iq_vec.push(Complex::new((iq[0] as f32 - 127.5) / 127.5,
                                 (iq[1] as f32 - 127.5) / 127.5));
    }

    Some(iq_vec)
}

fn demod_fm(iq: Vec<Complex<f32>>, gain: f32) -> VecDeque<f32> {
    let mut demod_queue: VecDeque<f32> = VecDeque::with_capacity(iq.len());
    let mut prev = iq[0];
    for ix in 0..iq.len() {
        let c: Complex<f32> = prev.conj() * iq[ix];
        demod_queue.push_back(c.re.atan2(c.im) * gain);
        prev = iq[ix];
    }
    demod_queue
}
<commit_msg>Added an offset from the center frequency to avoid DC spike<commit_after>extern crate num;
extern crate portaudio;
extern crate rtlsdr;

use portaudio as pa;
use rtlsdr::{RTLSDRDevice, RTLSDRError};
use num::complex::Complex;
use std::env;
use std::collections::VecDeque;
use std::f32::consts::PI;
use std::error::Error;

// Threading libraries.
use std::thread;
use std::sync::mpsc::channel;

// Make sure that the buffer size is radix-2, otherwise the read_sync function
// will fail with an error code of -8.
const BUF_SIZE: usize = 131072;

// Most other sample rates fail, but this one works for my particular device.
// I will investigate exactly what's happening here and generate a list of
// possible sample rates (if there are others besides this one).
const SAMPLE_RATE: u32 = 1140000;

// Here is the sample rate of the output waveform we'll try to use.
const SAMPLE_RATE_AUDIO: f64 = 48000.0;

const DEVIATION: u32 = 2500;

const F_OFFSET: u32 = 250000;

const BANDWIDTH: u32 = 200000;

const CHANNELS: i32 = 2;

fn main() {
    match run() {
        Ok(()) => (),
        Err(e) => println!("ERROR: {}", e),
    }
}

struct RTLSDR {
    rtlsdr: RTLSDRDevice,
}

unsafe impl Send for RTLSDR {}

fn run() -> Result<(), Box<Error>> {
    let sdr_index: i32 = env::args().nth(1)
        .expect("Please specify the SDR index")
        .parse::<i32>()?;
    let fm_freq_mhz: f32 = env::args().nth(2)
        .expect("Please specify a center frequency (MHz).")
        .parse::<f32>()?;

    // Initialize the PortAudio class.
    let audio = try!(pa::PortAudio::new());
    let settings =
        try!(audio.default_output_stream_settings(CHANNELS,
                                                  SAMPLE_RATE_AUDIO,
                                                  BUF_SIZE as u32));
    let mut stream = try!(audio.open_blocking_stream(settings));
    try!(stream.start());

    let fm_freq: u32 = (fm_freq_mhz * 1e6) as u32 - F_OFFSET;
    let mut sdr = RTLSDR { rtlsdr: init_sdr(sdr_index, fm_freq).unwrap() };

    let dec_rate = (SAMPLE_RATE as f64 / BANDWIDTH as f64) as usize;
    let (tx, rx) = channel();

    thread::spawn(move || {
        loop {
            // Read the samples and decimate down to match the sample rate of
            // the audio that'll be going out.
            let tx = tx.clone();
            let bytes = sdr.rtlsdr.read_sync(BUF_SIZE).unwrap();
            let iq_vec = decimate(read_samples(bytes, SAMPLE_RATE, F_OFFSET)
                                      .unwrap()
                                      .as_slice(),
                                  dec_rate);

            // After decimation, demodulate the signal.
            //
            // TODO: from what I've read there should be some decimation here as
            // well but with narrowband FM, we decimate the signal into
            // nothingness. Need to figure out what we do between wideband and
            // narrowband FM.
            let demod_iq = demod_fm(iq_vec);
            tx.send(demod_iq).unwrap();
        }
    });

    loop {
        // Get the write stream and write our samples to the stream.
        let mut demod_iq = rx.recv().unwrap();
        let out_frames = match stream.write_available() {
            Ok(available) => {
                match available {
                    pa::StreamAvailable::Frames(frames) => frames as u32,
                    _ => return Err(Box::new(pa::error::Error::NullCallback)),
                }
            }
            Err(e) => return Err(Box::new(e)),
        };

        let buffer_frames = (demod_iq.len() / CHANNELS as usize) as u32;
        let write_frames = if buffer_frames >= out_frames {
            out_frames
        } else {
            buffer_frames
        };
        let n_write_samples = write_frames as usize * CHANNELS as usize;

        stream.write(write_frames,
                   |output| for ix in 0..n_write_samples as usize {
                       output[ix] = demod_iq.pop_front().unwrap();
                   })?;
    }
}

fn init_sdr(sdr_index: i32, fm_freq: u32) -> Result<RTLSDRDevice, RTLSDRError> {
    let mut sdr = try!(rtlsdr::open(sdr_index));
    try!(sdr.set_center_freq(fm_freq));
    try!(sdr.set_sample_rate(SAMPLE_RATE));
    try!(sdr.set_tuner_bandwidth(BANDWIDTH));
    try!(sdr.set_agc_mode(true));
    try!(sdr.reset_buffer());
    Ok(sdr)
}

fn decimate_queue<T: Copy>(signal: VecDeque<T>,
                           dec_rate: usize)
                           -> VecDeque<T> {
    let mut ix = 0;
    let mut signal_dec = VecDeque::<T>::new();
    while ix < signal.len() {
        signal_dec.push_back(*(signal.get(ix).unwrap()));
        ix += dec_rate;
    }
    signal_dec
}

fn decimate(signal: &[Complex<f32>], dec_rate: usize) -> Vec<Complex<f32>> {
    let mut ix = 0;
    let new_size = (signal.len() / dec_rate + 1) as usize;
    let mut signal_dec = Vec::<Complex<f32>>::with_capacity(new_size);
    while ix < signal.len() {
        signal_dec.push(signal[ix]);
        ix += dec_rate;
    }
    signal_dec
}

fn read_samples(bytes: Vec<u8>,
                sample_freq: u32,
                offset: u32)
                -> Option<Vec<Complex<f32>>> {
    let j: Complex<f32> = Complex::i();
    let mut fc_ds: Complex<f32> = (-1.0 * j * 2.0 * PI * offset as f32 /
                                   sample_freq as f32) as
                                  Complex<f32>;
    fc_ds = fc_ds.exp();

    // First, check that we've been given an even number of bytes. Not sure
    // what to do if we don't get I and Q.
    let bytes_len = bytes.len();
    if bytes_len % 2 != 0 {
        println!("Received an odd number of samples.");
        return None;

    }
    let mut iq_vec: Vec<Complex<f32>> = Vec::with_capacity(bytes_len / 2);

    // Write the values to the complex value and normalize from [0, 255] to
    // [-1, 1]. If we don't normalize, it seems we don't get anything.
    for iq in bytes.chunks(2) {
        let iq_cmplx = Complex::new((iq[0] as f32 - 127.0) / 127.0,
                                    (iq[1] as f32 - 127.0) / 127.0) *
                       fc_ds;
        iq_vec.push(iq_cmplx);
    }
    Some(iq_vec)
}

fn demod_fm(iq: Vec<Complex<f32>>) -> VecDeque<f32> {
    let mut demod_queue: VecDeque<f32> = VecDeque::with_capacity(iq.len());
    let mut prev = iq[0];
    for ix in 0..iq.len() {
        let c: Complex<f32> = prev.conj() * iq[ix];
        demod_queue.push_back(c.re.atan2(c.im));
        prev = iq[ix];
    }
    demod_queue
}
<|endoftext|>


<commit_before>extern crate getopts;
extern crate pnet;
extern crate time;

use getopts::Options;
use std::os;
use std::net::{IpAddr};
use std::str::FromStr;

use pnet::transport::TransportProtocol::{Ipv4};
use pnet::old_packet::ip::IpNextHeaderProtocols;
use pnet::transport::TransportChannelType::{Layer3};
use pnet::transport::{transport_channel};
use pnet::old_packet::{Packet};
use pnet::old_packet::ipv4::{Ipv4Header, Ipv4Packet};

fn main() {
    let args: Vec<String> = os::args();
    let mut opts = Options::new();
    let matches = match opts.parse(args.tail()) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string())
    };
    match matches.free.len() {
        0 => panic!("must provide ip address"),
        _ => {}
    }
    let addr_arg = matches.free[0].as_slice();
    let addr: IpAddr = FromStr::from_str(addr_arg.as_slice()).unwrap();

    let protocol = Layer3(IpNextHeaderProtocols::Test1);
    let (mut icmp_sender, mut icmp_receiver) = match transport_channel(4096, protocol) {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("An error occurred when creating the transport channel:
                        {}", e)
    };
    let packet_vec: Vec<u8> = vec![];
    let packet = build_icmp_time_request_packet(packet_vec.as_slice());
}

fn build_icmp_time_request_packet<'p>(packet_slice: &'p [u8]) -> IcmpRequestPacket<'p> {
    let mut packet = IcmpRequestPacket::new(packet_slice);
    packet
}

pub trait IcmpPacket : Ipv4Packet {
    fn calculate_icmp_checksum(&self) -> u16 {
        let start_offset = self.get_header_length() as usize * 4;
        let chunked: Vec<u16> = self.packet().slice_from(start_offset).
            chunks(2).
            map(|x| x.iter().fold(0u16, |sum, y|
                                  if sum == 0 {
                                      sum + (*y as u16) << 4
                                  }
                                  else {
                                      sum + *y as u16
                                  }
                                 )).collect();
        ones_complement_sum(chunked.as_slice())
    }
}

struct IcmpRequestPacket<'p> {
    ip_type: u8,
    ip_code: u8,
    ip_checksum: u16,
    identifier: u16,
    sequence: u16,
    originate_timestamp: u32,
    receive_timestamp: u32,
    transmit_timestamp: u32,
    packet: &'p [u8]
}

// pretty much copy-pasted from pnet, not sure if inline directives are necessary
impl<'a> Packet for IcmpRequestPacket<'a> {
    #[inline(always)]
    fn packet<'p>(&'p self) -> &'p [u8] { self.packet }

    #[inline(always)]
    fn payload<'p>(&'p self) -> &'p [u8] { &self.packet[20..] }
    // use of 20 here should be fine for our case, IHL is 5, no IP options in header
}

impl<'p> Ipv4Packet for IcmpRequestPacket<'p> {}

impl<'p> IcmpRequestPacket<'p> {
    fn new(packet_slice: &'p [u8]) -> IcmpRequestPacket {
        let mut packet = IcmpRequestPacket {
            ip_type: 13,
            ip_code: 0,
            ip_checksum: 0,
            identifier: 0,
            sequence: 0,
            originate_timestamp: 0,
            receive_timestamp: 0,
            transmit_timestamp: 0,
            packet: packet_slice
        };
        packet.set_originate_timestamp();
        // note we actually want to do this just before sending, not yet
        packet.set_checksum();
        packet
    }

    fn calculate_checksum(&self) -> u16 {
        // implement me
        let mut vec: Vec<u16> = vec![];
        vec.push((self.ip_type + self.ip_code) as u16);
        vec.push(self.ip_checksum);
        vec.push(self.identifier);
        vec.push(self.sequence);
        vec.push((self.originate_timestamp >> 16) as u16);
        vec.push(self.originate_timestamp as u16);
        vec.push((self.receive_timestamp >> 16) as u16);
        vec.push(self.receive_timestamp as u16);
        vec.push((self.transmit_timestamp >> 16) as u16);
        vec.push(self.transmit_timestamp as u16);
        ones_complement_sum(vec.as_slice())
    }

    fn set_checksum(&mut self) {
        self.ip_checksum = self.calculate_checksum();
    }

    fn set_originate_timestamp(&mut self) {
        self.originate_timestamp = time_after_utc();
    }
}

fn time_after_utc() -> u32 {
    let t = time::now_utc();
    let s: i32 = t.tm_hour * 3600 + t.tm_min * 60 + t.tm_sec;
    s as u32
}

fn ones_complement_sum(sl: &[u16]) -> u16 {
    let sum = sl.iter().fold(0u32, |sum, x| sum + (*x as u32));
    65535 - (sum as u16) - (sum >> 16) as u16
}

mod tests {
    use super::ones_complement_sum;

    // examples for ones complement sum taken from web
    #[test]
    fn test_ones_complement_sum() {
        let mut vec: Vec<u16> = vec![];
        vec.push(0x4500);
        vec.push(0x003c);
        vec.push(0x1c46);
        vec.push(0x4000);
        vec.push(0x4006);
        vec.push(0x0000);
        vec.push(0xac10);
        vec.push(0x0a63);
        vec.push(0xac10);
        vec.push(0x0a0c);
        let r = vec.as_slice();
        assert_eq!(ones_complement_sum(r), 0xB1E6);
    }

    #[test]
    fn another_test_ones_complement_sum() {
        let mut vec: Vec<u16> = vec![];
        vec.push(0x0800);
        vec.push(0x0000);
        vec.push(0x0001);
        vec.push(0x1008);
        vec.push(0x6162);
        vec.push(0x6364);
        vec.push(0x6566);
        vec.push(0x6768);
        vec.push(0x696a);
        vec.push(0x6b6c);
        vec.push(0x6d6e);
        vec.push(0x6f70);
        vec.push(0x7172);
        vec.push(0x7374);
        vec.push(0x7576);
        vec.push(0x7761);
        vec.push(0x6263);
        vec.push(0x6465);
        vec.push(0x6667);
        vec.push(0x6869);
        let r = vec.as_slice();
        assert_eq!(ones_complement_sum(r), 15699);
    }
}
<commit_msg>wip on MutIcmpRequestPacket, moving logic into IcmpPacket trait. more in line with how libpnet works<commit_after>extern crate getopts;
extern crate pnet;
extern crate time;

use getopts::Options;
use std::os;
use std::net::{IpAddr};
use std::str::FromStr;

use pnet::transport::TransportProtocol::{Ipv4};
use pnet::old_packet::ip::IpNextHeaderProtocols;
use pnet::transport::TransportChannelType::{Layer3};
use pnet::transport::{transport_channel};
use pnet::old_packet::{Packet};
use pnet::old_packet::ipv4::{Ipv4Header, Ipv4Packet};

fn main() {
    let args: Vec<String> = os::args();
    let mut opts = Options::new();
    let matches = match opts.parse(args.tail()) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string())
    };
    match matches.free.len() {
        0 => panic!("must provide ip address"),
        _ => {}
    }
    let addr_arg = matches.free[0].as_slice();
    let addr: IpAddr = FromStr::from_str(addr_arg.as_slice()).unwrap();

    let protocol = Layer3(IpNextHeaderProtocols::Test1);
    let (mut icmp_sender, mut icmp_receiver) = match transport_channel(4096, protocol) {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("An error occurred when creating the transport channel:
                        {}", e)
    };
    let packet_vec: Vec<u8> = vec![];
    let packet = build_icmp_time_request_packet(packet_vec.as_slice());
}

fn build_icmp_time_request_packet<'p>(packet_slice: &'p [u8]) -> IcmpRequestPacket<'p> {
    let mut packet = IcmpRequestPacket::new(packet_slice);
    packet
}

pub trait IcmpPacket : Ipv4Packet {
    fn calculate_icmp_checksum(&self) -> u16 {
        // start at end of IP header packet
        let start_offset = self.get_header_length() as usize * 4;
        let chunked: Vec<u16> = self.packet().slice_from(start_offset).
            chunks(2).
            map(|x| x.iter().fold(0u16, |sum, y|
                                  if sum == 0 {
                                      sum + (*y as u16) << 4
                                  }
                                  else {
                                      sum + *y as u16
                                  }
                                 )).collect();
        ones_complement_sum(chunked.as_slice())
    }

    fn start_of_icmp(&self) -> usize {
        self.get_header_length() as usize * 4
    }

    fn get_ip_type(&self) -> u8 {
        let start = self.start_of_icmp();
        self.packet()[start]
    }

    fn get_ip_code(&self) -> u8 {
        let start = self.start_of_icmp();
        self.packet()[start + 1]
    }

    fn get_ip_checksum(&self) -> u16 {
        let start = self.start_of_icmp();
        let b1 = (self.packet()[start + 2] as u16) << 8;
        let b2 = (self.packet()[start + 3] as u16);
        b1 | b2
    }

    fn get_identifier(&self) -> u16 {
        let start = self.start_of_icmp();
        let b1 = (self.packet()[start + 4] as u16) << 8;
        let b2 = (self.packet()[start + 5] as u16);
        b1 | b2
    }

    fn get_sequence(&self) -> u16 {
        let start = self.start_of_icmp();
        let b1 = (self.packet()[start + 6] as u16) << 8;
        let b2 = (self.packet()[start + 7] as u16);
        b1 | b2
    }

    fn get_originate_timestamp(&self) -> u32 {
        let start = self.start_of_icmp();
        let b1 = (self.packet()[start + 8] as u32) << 24;
        let b2 = (self.packet()[start + 9] as u32) << 16;
        let b3 = (self.packet()[start + 10] as u32) << 8;
        let b4 = (self.packet()[start + 11] as u32);
        b1 | b2 | b3 | b4
    }

    fn get_receive_timestamp(&self) -> u32 {
        let start = self.start_of_icmp();
        let b1 = (self.packet()[start + 12] as u32) << 24;
        let b2 = (self.packet()[start + 13] as u32) << 16;
        let b3 = (self.packet()[start + 14] as u32) << 8;
        let b4 = (self.packet()[start + 15] as u32);
        b1 | b2 | b3 | b4
    }

    fn get_transmit_timestamp(&self) -> u32 {
        let start = self.start_of_icmp();
        let b1 = (self.packet()[start + 16] as u32) << 24;
        let b2 = (self.packet()[start + 17] as u32) << 16;
        let b3 = (self.packet()[start + 18] as u32) << 8;
        let b4 = (self.packet()[start + 19] as u32);
        b1 | b2 | b3 | b4
    }
}

pub struct IcmpRequestPacket<'p> {
    packet: &'p [u8]
}

pub struct MutIcmpRequestPacket<'p> {
    packet: &'p mut [u8]
}

// pretty much copy-pasted from pnet, not sure if inline directives are necessary
impl<'a> Packet for IcmpRequestPacket<'a> {
    #[inline(always)]
    fn packet<'p>(&'p self) -> &'p [u8] { self.packet }

    #[inline(always)]
    fn payload<'p>(&'p self) -> &'p [u8] { &self.packet[20..] }
    // use of 20 here should be fine for our case, IHL is 5, no IP options in header
}

impl<'a> Packet for MutIcmpRequestPacket<'a> {
    #[inline(always)]
    fn packet<'p>(&'p self) -> &'p [u8] { self.packet }

    #[inline(always)]
    fn payload<'p>(&'p self) -> &'p [u8] { &self.packet[20..] }
    // use of 20 here should be fine for our case, IHL is 5, no IP options in header
}

impl<'p> Ipv4Packet for IcmpRequestPacket<'p> {}
impl<'p> IcmpPacket for IcmpRequestPacket<'p> {}

impl<'p> Ipv4Packet for MutIcmpRequestPacket<'p> {}
impl<'p> IcmpPacket for MutIcmpRequestPacket<'p> {}

impl<'p> IcmpRequestPacket<'p> {
    fn new(packet_slice: &'p [u8]) -> IcmpRequestPacket {
        IcmpRequestPacket {
            packet: packet_slice
        }
    }
}

impl<'p> MutIcmpRequestPacket<'p> {
    fn new(packet_slice: &'p mut [u8]) -> MutIcmpRequestPacket {
        MutIcmpRequestPacket {
            packet: packet_slice
        }
    }

    fn set_checksum(&mut self) {
        let start = self.start_of_icmp();
        let checksum = self.calculate_icmp_checksum();
        self.packet[start + 2] = (checksum >> 8) as u8;
        self.packet[start + 3] = (checksum & 0xf) as u8;
    }

    fn set_originate_timestamp(&mut self) {
        let start = self.start_of_icmp();
        let time = time_after_utc();
        // dbl check this offset
        self.packet[start + 8] = (time >> 24) as u8;
        self.packet[start + 9] = ((time >> 16) & 0xf0) as u8;
        self.packet[start + 10] = ((time >> 8) & 0xff0) as u8;
        self.packet[start + 11] = time as u8;
    }
}

fn time_after_utc() -> u32 {
    let t = time::now_utc();
    let s: i32 = t.tm_hour * 3600 + t.tm_min * 60 + t.tm_sec;
    s as u32
}

fn ones_complement_sum(sl: &[u16]) -> u16 {
    let sum = sl.iter().fold(0u32, |sum, x| sum + (*x as u32));
    65535 - (sum as u16) - (sum >> 16) as u16
}

mod tests {
    use super::ones_complement_sum;

    // examples for ones complement sum taken from web
    #[test]
    fn test_ones_complement_sum() {
        let mut vec: Vec<u16> = vec![];
        vec.push(0x4500);
        vec.push(0x003c);
        vec.push(0x1c46);
        vec.push(0x4000);
        vec.push(0x4006);
        vec.push(0x0000);
        vec.push(0xac10);
        vec.push(0x0a63);
        vec.push(0xac10);
        vec.push(0x0a0c);
        let r = vec.as_slice();
        assert_eq!(ones_complement_sum(r), 0xB1E6);
    }

    #[test]
    fn another_test_ones_complement_sum() {
        let mut vec: Vec<u16> = vec![];
        vec.push(0x0800);
        vec.push(0x0000);
        vec.push(0x0001);
        vec.push(0x1008);
        vec.push(0x6162);
        vec.push(0x6364);
        vec.push(0x6566);
        vec.push(0x6768);
        vec.push(0x696a);
        vec.push(0x6b6c);
        vec.push(0x6d6e);
        vec.push(0x6f70);
        vec.push(0x7172);
        vec.push(0x7374);
        vec.push(0x7576);
        vec.push(0x7761);
        vec.push(0x6263);
        vec.push(0x6465);
        vec.push(0x6667);
        vec.push(0x6869);
        let r = vec.as_slice();
        assert_eq!(ones_complement_sum(r), 15699);
    }
}
<|endoftext|>