<commit_before>#[macro_use]
extern crate clap;
extern crate chrono;
#[macro_use]
extern crate serde_derive;
extern crate serde;
extern crate serde_json;
extern crate colored;

use std::env;
use std::fs;
use std::fs::{OpenOptions, File, Metadata};
use std::io::{Write, ErrorKind};
use std::path::PathBuf;
use std::collections::{HashMap};

use chrono::{DateTime, Date, Utc, Local, Duration};
use clap::{App, Arg, AppSettings, SubCommand};
use serde_json::Error;
use colored::*;

#[derive(Serialize, Deserialize, Debug, Clone)]
struct Period {
    project: String,
    start_time: DateTime<Utc>,
    end_time: Option<DateTime<Utc>>,
}

fn main() {
    let matches = App::new("Doug")
        .version(crate_version!())
        .about("A time tracking command-line utility")
        .author(crate_authors!())
        .settings(&[
            AppSettings::DeriveDisplayOrder,
            AppSettings::GlobalVersion,
            AppSettings::SubcommandRequiredElseHelp,
            AppSettings::VersionlessSubcommands,
            AppSettings::DisableHelpSubcommand,
        ])
        .subcommand(SubCommand::with_name("start")
            .about("Track new or existing project")
            .arg(Arg::with_name("project")
                .help("project to track")
                .required(true)))
        .subcommand(SubCommand::with_name("status")
            .about("Display elapsed time, start time, and running project name"))
        .subcommand(SubCommand::with_name("stop")
            .about("Stop any running projects"))
        .subcommand(SubCommand::with_name("cancel")
            .about("Stop running project and remove most recent time interval"))
        .subcommand(SubCommand::with_name("restart")
            .about("Track last running project"))
        .subcommand(SubCommand::with_name("log")
            .about("Display time intervals across all projects"))
        .subcommand(SubCommand::with_name("report")
            .about("Display aggregate time from projects"))
        .subcommand(SubCommand::with_name("amend")
            .about("Change name of currently running project")
            .arg(Arg::with_name("project")
                .help("new project name")
                .required(true)))
        .subcommand(SubCommand::with_name("edit")
            .about("Edit last frame or currently running frame")
            .arg(Arg::with_name("project")
                .help("project to track")))
        .subcommand(SubCommand::with_name("delete")
            .about("Delete all intervals for project")
            .arg(Arg::with_name("project")
                .help("new project name")
                .required(true)))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("start") {
        if matches.is_present("project") {
            start(matches.value_of("project").unwrap());
        }
    }

    if let Some(matches) = matches.subcommand_matches("amend") {
        if matches.is_present("project") {
            amend(matches.value_of("project").expect("missing project name"));
        }
    }

    if let Some(matches) = matches.subcommand_matches("edit") {
        edit(matches.value_of("project"));
    }

    if let Some(matches) = matches.subcommand_matches("delete") {
        delete(matches.value_of("project").expect("missing project name"));
    }

    match matches.subcommand_name() {
        Some("status") => status(),
        Some("stop") => stop(),
        Some("cancel") => cancel(),
        Some("restart") => restart(),
        Some("log") => log(),
        Some("report") => report(),
        _ => {},
    }
}

fn start(project_name: &str) {
    let mut periods = get_periods();

    if !periods.is_empty() {
        let last_index = periods.len() - 1;
        if let Some(period) = periods.get_mut(last_index) {
            if period.end_time.is_none() {
                let message = format!("project {} is being tracked", period.project);
                eprintln!("Error: {}",  message.red());
                return eprintln!("Try stopping your current project with {} first.", "stop".blue());
            }
        }
    }
    let current_period = create_period(project_name);
    // store current period in file
    print!("Started tracking project {} at {}", current_period.project.blue(), humanize_time(current_period.start_time));
    periods.push(current_period);
    save_periods(periods.to_vec());
}

fn status() {
    let mut periods = get_periods();
    if let Some(period) = periods.pop() {
        if period.end_time.is_none() {
            let diff = Utc::now().signed_duration_since(period.start_time);
            return println!("Project {} started {} ({})", period.project.magenta(), humanize_duration(diff), humanize_datetime(period.start_time).blue());
        }
    }
    eprintln!("No running project");
}

fn stop() {
    let mut periods = get_periods();
    let mut updated_period = false;
    let last_index = periods.len() - 1;
    if periods.len() > 0 {
        if let Some(period) = periods.get_mut(last_index) {
            if period.end_time.is_none() {
                period.end_time = Some(Utc::now());
                updated_period = true;
                let diff = Utc::now().signed_duration_since(period.start_time);
                println!("Stopped project {}, started {}", period.project.blue(), humanize_duration(diff));
            }
        }
    }
    if updated_period {
        save_periods(periods);
    } else {
        eprintln!("Error: {}", "No project started.".red());
    }
}

fn cancel() {
    let mut periods = get_periods();
    if let Some(period) = periods.pop() {
        if period.end_time.is_none() {
            save_periods(periods);
            let diff = Utc::now().signed_duration_since(period.start_time);
            println!("Canceled project {}, started {}", period.project.blue(), humanize_duration(diff));
            return
        }
    }
    eprintln!("Error: {}", "No project started".red());
}

fn restart() {
    let periods = get_periods();
    let mut new_periods = periods.to_vec();
    if let Some(period) = periods.last() {
        if !period.end_time.is_none() {
            let new_period = create_period(&period.project);
            new_periods.push(new_period);
            save_periods(new_periods);
            return println!("Tracking last running project: {}", period.project.blue());
        } else {
            let message = format!("No project to restart. Project {} is being tracked", period.project);
            eprintln!("Error: {}",  message.red());
            return eprintln!("Try stopping your current project with {} first.", "stop".blue());
        }
    }
    eprintln!("Error: {}", "No previous project to restart".red());
}

fn log() {
    let periods = get_periods();
    let mut days: HashMap<Date<chrono::Local>, Vec<Period>> = HashMap::new();

    for period in periods.iter() {
        let time = period.start_time.with_timezone(&Local).date();
        days.entry(time).or_insert(Vec::new()).push(period.clone());
    }
    for (date, day) in days.iter() {
        let d = day.into_iter().fold(Duration::zero(), |acc, ref x| acc + (x.end_time.unwrap_or(Utc::now()).signed_duration_since(x.start_time)));
        println!("{date} ({duration})", date=date.with_timezone(&Local).format("%A %-d %B %Y").to_string().green(), duration=format_duration(d).bold());
        for period in day.iter() {
            match period.end_time {
                Some(end_time) => {
                    let diff = end_time.signed_duration_since(period.start_time);
                    println!("    {start} to {end} {project} {duration}",
                        project=period.project.blue(),
                        start=humanize_time(period.start_time),
                        end=humanize_time(end_time),
                        duration=format_duration(diff));
                },
                None => {
                    let diff = Utc::now().signed_duration_since(period.start_time);
                    println!("    {start} to --:-- {project} {duration}",
                        project=period.project.blue(),
                        start=humanize_time(period.start_time),
                        duration=format_duration(diff));
                },
            }
        }
    }
}

fn report() {
    let periods = get_periods();
    let mut days: HashMap<String, Vec<Period>> = HashMap::new();
    let mut start_date = Utc::now().with_timezone(&Local).date();
    let mut results = Vec::new();

    for period in periods.iter() {
        days.entry(period.project.clone()).or_insert(Vec::new()).push(period.clone());
    }
    for (project, intervals) in days.iter() {
        let d = intervals.into_iter().fold(Duration::zero(), |acc, ref x| acc + (x.end_time.unwrap_or(Utc::now()).signed_duration_since(x.start_time)));
        for x in intervals.iter() {
            if x.start_time.with_timezone(&Local).date() < start_date {
                start_date = x.start_time.with_timezone(&Local).date();
            }
        }
        results.push(format!("{project} {duration}", project=project.green(), duration=format_duration(d).bold()));
    }
    println!("{start} -> {end}",
        start=start_date.format("%A %-d %B %Y").to_string().blue(),
        end=Utc::now().format("%A %-d %B %Y").to_string().blue());
    for x in results.iter() {
        println!("{}", x);
    }
}

fn amend(project_name: &str) {
    println!("Amend project: {:?}", project_name);
    unimplemented!();
}

fn edit(project_name: Option<&str>) {
    if let Some(name) = project_name {
        println!("Edit project: {}", name);
    }
    unimplemented!();
}

fn delete(project_name: &str) {
    println!("Delete project: {:?}", project_name);
    unimplemented!();
}

fn humanize_datetime(time: DateTime<Utc>) -> String {
    time.with_timezone(&Local).format("%F %H:%M").to_string()
}

fn humanize_time(time: DateTime<Utc>) -> String {
    time.with_timezone(&Local).format("%H:%M").to_string()
}

fn format_duration(duration: Duration) -> String{
    let days = duration.num_days();
    let hours = duration.num_hours();
    let minutes = duration.num_minutes();
    let seconds = duration.num_seconds();
    if minutes == 0 {
        return format!("{}s", seconds)
    } else if hours == 0 {
        return format!("{minutes}m", minutes=minutes)
    } else if days == 0 {
        return format!("{hours}h, {minutes}m", hours=hours, minutes=minutes/(hours*60))
    } else {
        return format!("{days}d, {hours}h, {minutes}m", hours=hours/(days*24), minutes=minutes/(hours*60), days=days)
    }
}

fn humanize_duration(time: Duration) -> String {
    let hours = time.num_hours();
    let minutes = time.num_minutes();
    let seconds = time.num_seconds();
    if minutes == 0 {
        if seconds < 5 {
            return String::from("just now")
        }
        return String::from("seconds ago")
    } else if hours == 0 {
        if minutes == 1 {
            return format!("{minutes} minute ago", minutes=minutes)
        } else {
            return format!("{minutes} minutes ago", minutes=minutes)
        }
    } else {
        return format!("{hours}:{minutes}", hours=hours, minutes=minutes)
    }
}

fn create_period(project: &str) -> Period {
    Period {
        project: String::from(project),
        start_time: Utc::now(),
        end_time: None,
    }
}

fn get_config_folder() -> PathBuf {
    let home_dir = env::var("HOME").expect("Failed to find home directory from environment 'HOME'");
    let mut config_folder = PathBuf::from(home_dir);
    config_folder.push(".config/doug");
    config_folder
}

fn get_config_file_path() -> PathBuf {
    let config_folder = get_config_folder();
    let mut config_file = PathBuf::from(&config_folder);
    config_file.push("periods.json");
    config_file
}

fn get_config_back_file_path() -> PathBuf {
    let config_file = get_config_file_path();
    config_file.with_extension("json-backup")
}

fn get_config_file() -> (File, Metadata) {
    let config_file = get_config_file_path();
    (OpenOptions::new()
        .create(true)
        .read(true)
        .write(true)
        .open(&config_file)
        .expect(&format!("Couldn't open datafile: {:?}", config_file)),
    fs::metadata(&config_file).expect("Couldn't access datafile metadata")
    )
}

fn get_periods() -> Vec<Period> {
    let config_folder = get_config_folder();

    match fs::create_dir(&config_folder) {
        Err(ref error) if error.kind() == ErrorKind::AlreadyExists => {},
        Err(error) => panic!("There was a problem creating the config directory: {:?}", error),
        Ok(_) => {},
    }

    let (file, metadata) = get_config_file();

    let periods: Result<Vec<Period>, Error> = serde_json::from_reader(file);
    let periods = match periods {
        Ok(p) => p,
        Err(_) if metadata.len() > 0 => Vec::new(),
        Err(error) => panic!("There was a serialization issue: {:?}", error),
    };
    periods
}

fn save_periods(periods: Vec<Period>) {
    let serialized = serde_json::to_string(&periods).expect("Couldn't serialize data to string");
    let config_file = get_config_file_path();
    let config_file_backup = get_config_back_file_path();
    fs::copy(&config_file, &config_file_backup).expect("Couldn't create backup file");
    let mut file = OpenOptions::new()
                    .create(true)
                    .write(true)
                    .truncate(true)
                    .open(&config_file)
                    .expect("Couldn't open file for saving period.");

    file.write_all(serialized.as_bytes()).expect("Couldn't write serialized data to file");
}
<commit_msg>Add delete subcommand<commit_after>#[macro_use]
extern crate clap;
extern crate chrono;
#[macro_use]
extern crate serde_derive;
extern crate serde;
extern crate serde_json;
extern crate colored;

use std::env;
use std::fs;
use std::fs::{OpenOptions, File, Metadata};
use std::io::{Write, ErrorKind};
use std::path::PathBuf;
use std::collections::{HashMap};

use chrono::{DateTime, Date, Utc, Local, Duration};
use clap::{App, Arg, AppSettings, SubCommand};
use serde_json::Error;
use colored::*;

#[derive(Eq, PartialEq, Serialize, Deserialize, Debug, Clone)]
struct Period {
    project: String,
    start_time: DateTime<Utc>,
    end_time: Option<DateTime<Utc>>,
}

fn main() {
    let matches = App::new("Doug")
        .version(crate_version!())
        .about("A time tracking command-line utility")
        .author(crate_authors!())
        .settings(&[
            AppSettings::DeriveDisplayOrder,
            AppSettings::GlobalVersion,
            AppSettings::SubcommandRequiredElseHelp,
            AppSettings::VersionlessSubcommands,
            AppSettings::DisableHelpSubcommand,
        ])
        .subcommand(SubCommand::with_name("start")
            .about("Track new or existing project")
            .arg(Arg::with_name("project")
                .help("project to track")
                .required(true)))
        .subcommand(SubCommand::with_name("status")
            .about("Display elapsed time, start time, and running project name"))
        .subcommand(SubCommand::with_name("stop")
            .about("Stop any running projects"))
        .subcommand(SubCommand::with_name("cancel")
            .about("Stop running project and remove most recent time interval"))
        .subcommand(SubCommand::with_name("restart")
            .about("Track last running project"))
        .subcommand(SubCommand::with_name("log")
            .about("Display time intervals across all projects"))
        .subcommand(SubCommand::with_name("report")
            .about("Display aggregate time from projects"))
        .subcommand(SubCommand::with_name("amend")
            .about("Change name of currently running project")
            .arg(Arg::with_name("project")
                .help("new project name")
                .required(true)))
        .subcommand(SubCommand::with_name("edit")
            .about("Edit last frame or currently running frame")
            .arg(Arg::with_name("project")
                .help("project to track")))
        .subcommand(SubCommand::with_name("delete")
            .about("Delete all intervals for project")
            .arg(Arg::with_name("project")
                .help("new project name")
                .required(true)))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("start") {
        if matches.is_present("project") {
            start(matches.value_of("project").unwrap());
        }
    }

    if let Some(matches) = matches.subcommand_matches("amend") {
        if matches.is_present("project") {
            amend(matches.value_of("project").expect("missing project name"));
        }
    }

    if let Some(matches) = matches.subcommand_matches("edit") {
        edit(matches.value_of("project"));
    }

    if let Some(matches) = matches.subcommand_matches("delete") {
        delete(matches.value_of("project").expect("missing project name"));
    }

    match matches.subcommand_name() {
        Some("status") => status(),
        Some("stop") => stop(),
        Some("cancel") => cancel(),
        Some("restart") => restart(),
        Some("log") => log(),
        Some("report") => report(),
        _ => {},
    }
}

fn start(project_name: &str) {
    let mut periods = get_periods();

    if !periods.is_empty() {
        let last_index = periods.len() - 1;
        if let Some(period) = periods.get_mut(last_index) {
            if period.end_time.is_none() {
                let message = format!("project {} is being tracked", period.project);
                eprintln!("Error: {}",  message.red());
                return eprintln!("Try stopping your current project with {} first.", "stop".blue());
            }
        }
    }
    let current_period = create_period(project_name);
    // store current period in file
    print!("Started tracking project {} at {}", current_period.project.blue(), humanize_time(current_period.start_time));
    periods.push(current_period);
    save_periods(periods.to_vec());
}

fn status() {
    let mut periods = get_periods();
    if let Some(period) = periods.pop() {
        if period.end_time.is_none() {
            let diff = Utc::now().signed_duration_since(period.start_time);
            return println!("Project {} started {} ({})", period.project.magenta(), humanize_duration(diff), humanize_datetime(period.start_time).blue());
        }
    }
    eprintln!("No running project");
}

fn stop() {
    let mut periods = get_periods();
    let mut updated_period = false;
    let last_index = periods.len() - 1;
    if periods.len() > 0 {
        if let Some(period) = periods.get_mut(last_index) {
            if period.end_time.is_none() {
                period.end_time = Some(Utc::now());
                updated_period = true;
                let diff = Utc::now().signed_duration_since(period.start_time);
                println!("Stopped project {}, started {}", period.project.blue(), humanize_duration(diff));
            }
        }
    }
    if updated_period {
        save_periods(periods);
    } else {
        eprintln!("Error: {}", "No project started.".red());
    }
}

fn cancel() {
    let mut periods = get_periods();
    if let Some(period) = periods.pop() {
        if period.end_time.is_none() {
            save_periods(periods);
            let diff = Utc::now().signed_duration_since(period.start_time);
            println!("Canceled project {}, started {}", period.project.blue(), humanize_duration(diff));
            return
        }
    }
    eprintln!("Error: {}", "No project started".red());
}

fn restart() {
    let periods = get_periods();
    let mut new_periods = periods.to_vec();
    if let Some(period) = periods.last() {
        if !period.end_time.is_none() {
            let new_period = create_period(&period.project);
            new_periods.push(new_period);
            save_periods(new_periods);
            return println!("Tracking last running project: {}", period.project.blue());
        } else {
            let message = format!("No project to restart. Project {} is being tracked", period.project);
            eprintln!("Error: {}",  message.red());
            return eprintln!("Try stopping your current project with {} first.", "stop".blue());
        }
    }
    eprintln!("Error: {}", "No previous project to restart".red());
}

fn log() {
    let periods = get_periods();
    let mut days: HashMap<Date<chrono::Local>, Vec<Period>> = HashMap::new();

    for period in periods.iter() {
        let time = period.start_time.with_timezone(&Local).date();
        days.entry(time).or_insert(Vec::new()).push(period.clone());
    }
    for (date, day) in days.iter() {
        let d = day.into_iter().fold(Duration::zero(), |acc, ref x| acc + (x.end_time.unwrap_or(Utc::now()).signed_duration_since(x.start_time)));
        println!("{date} ({duration})", date=date.with_timezone(&Local).format("%A %-d %B %Y").to_string().green(), duration=format_duration(d).bold());
        for period in day.iter() {
            match period.end_time {
                Some(end_time) => {
                    let diff = end_time.signed_duration_since(period.start_time);
                    println!("    {start} to {end} {project} {duration}",
                        project=period.project.blue(),
                        start=humanize_time(period.start_time),
                        end=humanize_time(end_time),
                        duration=format_duration(diff));
                },
                None => {
                    let diff = Utc::now().signed_duration_since(period.start_time);
                    println!("    {start} to --:-- {project} {duration}",
                        project=period.project.blue(),
                        start=humanize_time(period.start_time),
                        duration=format_duration(diff));
                },
            }
        }
    }
}

fn report() {
    let periods = get_periods();
    let mut days: HashMap<String, Vec<Period>> = HashMap::new();
    let mut start_date = Utc::now().with_timezone(&Local).date();
    let mut results = Vec::new();

    for period in periods.iter() {
        days.entry(period.project.clone()).or_insert(Vec::new()).push(period.clone());
    }
    for (project, intervals) in days.iter() {
        let d = intervals.into_iter().fold(Duration::zero(), |acc, ref x| acc + (x.end_time.unwrap_or(Utc::now()).signed_duration_since(x.start_time)));
        for x in intervals.iter() {
            if x.start_time.with_timezone(&Local).date() < start_date {
                start_date = x.start_time.with_timezone(&Local).date();
            }
        }
        results.push(format!("{project} {duration}", project=project.green(), duration=format_duration(d).bold()));
    }
    println!("{start} -> {end}",
        start=start_date.format("%A %-d %B %Y").to_string().blue(),
        end=Utc::now().format("%A %-d %B %Y").to_string().blue());
    for x in results.iter() {
        println!("{}", x);
    }
}

fn amend(project_name: &str) {
    println!("Amend project: {:?}", project_name);
    unimplemented!();
}

fn edit(project_name: Option<&str>) {
    if let Some(name) = project_name {
        println!("Edit project: {}", name);
    }
    unimplemented!();
}

fn delete(project_name: &str) {
    let mut periods = get_periods();
    let filtered_periods = periods.clone().into_iter().filter(|x| x.project != project_name).collect();
    if filtered_periods == periods {
        eprintln!("Error: {}", "Project not found.".red());
    } else {
        save_periods(filtered_periods);
        println!("Deleted project {project}", project=project_name.blue());
    }
}

fn humanize_datetime(time: DateTime<Utc>) -> String {
    time.with_timezone(&Local).format("%F %H:%M").to_string()
}

fn humanize_time(time: DateTime<Utc>) -> String {
    time.with_timezone(&Local).format("%H:%M").to_string()
}

fn format_duration(duration: Duration) -> String{
    let days = duration.num_days();
    let hours = duration.num_hours();
    let minutes = duration.num_minutes();
    let seconds = duration.num_seconds();
    if minutes == 0 {
        return format!("{}s", seconds)
    } else if hours == 0 {
        return format!("{minutes}m", minutes=minutes)
    } else if days == 0 {
        return format!("{hours}h, {minutes}m", hours=hours, minutes=minutes/(hours*60))
    } else {
        return format!("{days}d, {hours}h, {minutes}m", hours=hours/(days*24), minutes=minutes/(hours*60), days=days)
    }
}

fn humanize_duration(time: Duration) -> String {
    let hours = time.num_hours();
    let minutes = time.num_minutes();
    let seconds = time.num_seconds();
    if minutes == 0 {
        if seconds < 5 {
            return String::from("just now")
        }
        return String::from("seconds ago")
    } else if hours == 0 {
        if minutes == 1 {
            return format!("{minutes} minute ago", minutes=minutes)
        } else {
            return format!("{minutes} minutes ago", minutes=minutes)
        }
    } else {
        return format!("{hours}:{minutes}", hours=hours, minutes=minutes)
    }
}

fn create_period(project: &str) -> Period {
    Period {
        project: String::from(project),
        start_time: Utc::now(),
        end_time: None,
    }
}

fn get_config_folder() -> PathBuf {
    let home_dir = env::var("HOME").expect("Failed to find home directory from environment 'HOME'");
    let mut config_folder = PathBuf::from(home_dir);
    config_folder.push(".config/doug");
    config_folder
}

fn get_config_file_path() -> PathBuf {
    let config_folder = get_config_folder();
    let mut config_file = PathBuf::from(&config_folder);
    config_file.push("periods.json");
    config_file
}

fn get_config_back_file_path() -> PathBuf {
    let config_file = get_config_file_path();
    config_file.with_extension("json-backup")
}

fn get_config_file() -> (File, Metadata) {
    let config_file = get_config_file_path();
    (OpenOptions::new()
        .create(true)
        .read(true)
        .write(true)
        .open(&config_file)
        .expect(&format!("Couldn't open datafile: {:?}", config_file)),
    fs::metadata(&config_file).expect("Couldn't access datafile metadata")
    )
}

fn get_periods() -> Vec<Period> {
    let config_folder = get_config_folder();

    match fs::create_dir(&config_folder) {
        Err(ref error) if error.kind() == ErrorKind::AlreadyExists => {},
        Err(error) => panic!("There was a problem creating the config directory: {:?}", error),
        Ok(_) => {},
    }

    let (file, metadata) = get_config_file();

    let periods: Result<Vec<Period>, Error> = serde_json::from_reader(file);
    let periods = match periods {
        Ok(p) => p,
        Err(_) if metadata.len() > 0 => Vec::new(),
        Err(error) => panic!("There was a serialization issue: {:?}", error),
    };
    periods
}

fn save_periods(periods: Vec<Period>) {
    let serialized = serde_json::to_string(&periods).expect("Couldn't serialize data to string");
    let config_file = get_config_file_path();
    let config_file_backup = get_config_back_file_path();
    fs::copy(&config_file, &config_file_backup).expect("Couldn't create backup file");
    let mut file = OpenOptions::new()
                    .create(true)
                    .write(true)
                    .truncate(true)
                    .open(&config_file)
                    .expect("Couldn't open file for saving period.");

    file.write_all(serialized.as_bytes()).expect("Couldn't write serialized data to file");
}
<|endoftext|>


<commit_before>extern crate pnet;
extern crate clap;

pub mod cli;
pub mod config;

use std::thread;
use std::time::Duration;
use std::net::Ipv4Addr;

use pnet::datalink::{self, NetworkInterface};
use pnet::datalink::Channel;
use pnet::packet::ethernet::MutableEthernetPacket;
use pnet::packet::arp::MutableArpPacket;
use pnet::util::MacAddr;
use pnet::packet::ethernet::EtherTypes;
use pnet::packet::MutablePacket;
use pnet::packet::arp::{ArpHardwareTypes, ArpOperations};

use config::Config;
use cli::cli_main;

fn send_arp_reply_packet(gateway: Ipv4Addr, source_mac: MacAddr, target_ip: Ipv4Addr, target_mac: MacAddr) {
    let interfaces = datalink::interfaces();

    // Get en1 interface in osx.
    let interfaces_name_match = |iface: &NetworkInterface| iface.name == "en1";
    let interface = interfaces.into_iter().filter(interfaces_name_match).next().unwrap();

    let(mut tx, _) = match datalink::channel(&interface, Default::default()) {
        Ok(Channel::Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("Unknown channel type"),
        Err(e) => panic!("Error happened {}", e),
    };

    /// ethernet_packet = Ethernet {
    ///     destination: [0xff, 0xff, 0xff, 0xff, 0xff, 0xff],
    ///     source: [0x28, 0xef, 0xf9, 0x5f, 0x8e, 0x2b],
    ///     ethertype: [0x08, 0x06], // Arp(0x0806)
    ///     payload: arp_packet
    /// }
    ///
    /// arp_packet = Arp {
    ///     hardware_type: [0x00, 0x01],
    ///     protocol_type: [0x08, 0x00], // Ipv4(0x0800)
    ///     hw_addr_len: [0x06],
    ///     proto_addr_len: [0x04],
    ///     operation: [0x00, 0x02], // Reply(0x0002)
    ///     sender_hw_addr: [0x28, 0xef, 0xf9, 0x5f, 0x8e, 0x2b],
    ///     sender_proto_addr: [0xc0, 0xa8, 0x00, 0x66], // Ipv4(192.168.0.102)
    ///     target_hw_addr: [0xff, 0xff, 0xff, 0xff, 0xff, 0xff], // Broadcast
    ///     target_proto_addr: [0xc0, 0xa8, 0x00, 0x65], // Ipv4(192.168.0.101)
    ///     payload: [],
    /// }

    let mut buffer: &mut [u8] = &mut [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];

    let mut ethernet_packet = MutableEthernetPacket::new(&mut buffer).unwrap();
    ethernet_packet.set_destination(target_mac);
    ethernet_packet.set_source(source_mac);
    ethernet_packet.set_ethertype(EtherTypes::Arp);

    let mut arp_buffer: &mut [u8] = &mut [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                          0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                          0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];

    let mut arp_packet = MutableArpPacket::new(&mut arp_buffer).unwrap();

    arp_packet.set_hardware_type(ArpHardwareTypes::Ethernet);
    arp_packet.set_protocol_type(EtherTypes::Ipv4);
    arp_packet.set_hw_addr_len(6);
    arp_packet.set_proto_addr_len(4);
    arp_packet.set_operation(ArpOperations::Reply);
    arp_packet.set_sender_hw_addr(source_mac);
    arp_packet.set_sender_proto_addr(gateway);
    arp_packet.set_target_hw_addr(target_mac);
    arp_packet.set_target_proto_addr(target_ip);

    ethernet_packet.set_payload(arp_packet.packet_mut());

    tx.send_to(&ethernet_packet.to_immutable(), Some(interface));
}

fn main() {
    let config: Config = cli_main();
    let mut packet_count: i32 = 0;

    loop {
        send_arp_reply_packet(config.source_ip, config.source_mac, config.target_ip, config.target_mac);
        thread::sleep(Duration::new(5, 0));

        packet_count += 1;

        println!("Sent {} ARP reply packets.", packet_count);
    }
}
<commit_msg>Use readable variable name for ethernet buffer<commit_after>extern crate pnet;
extern crate clap;

pub mod cli;
pub mod config;

use std::thread;
use std::time::Duration;
use std::net::Ipv4Addr;

use pnet::datalink::{self, NetworkInterface};
use pnet::datalink::Channel;
use pnet::packet::ethernet::MutableEthernetPacket;
use pnet::packet::arp::MutableArpPacket;
use pnet::util::MacAddr;
use pnet::packet::ethernet::EtherTypes;
use pnet::packet::MutablePacket;
use pnet::packet::arp::{ArpHardwareTypes, ArpOperations};

use config::Config;
use cli::cli_main;

fn send_arp_reply_packet(gateway: Ipv4Addr, source_mac: MacAddr, target_ip: Ipv4Addr, target_mac: MacAddr) {
    let interfaces = datalink::interfaces();

    // Get en1 interface in osx.
    let interfaces_name_match = |iface: &NetworkInterface| iface.name == "en1";
    let interface = interfaces.into_iter().filter(interfaces_name_match).next().unwrap();

    let(mut tx, _) = match datalink::channel(&interface, Default::default()) {
        Ok(Channel::Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("Unknown channel type"),
        Err(e) => panic!("Error happened {}", e),
    };

    /// ethernet_packet = Ethernet {
    ///     destination: [0xff, 0xff, 0xff, 0xff, 0xff, 0xff],
    ///     source: [0x28, 0xef, 0xf9, 0x5f, 0x8e, 0x2b],
    ///     ethertype: [0x08, 0x06], // Arp(0x0806)
    ///     payload: arp_packet
    /// }
    ///
    /// arp_packet = Arp {
    ///     hardware_type: [0x00, 0x01],
    ///     protocol_type: [0x08, 0x00], // Ipv4(0x0800)
    ///     hw_addr_len: [0x06],
    ///     proto_addr_len: [0x04],
    ///     operation: [0x00, 0x02], // Reply(0x0002)
    ///     sender_hw_addr: [0x28, 0xef, 0xf9, 0x5f, 0x8e, 0x2b],
    ///     sender_proto_addr: [0xc0, 0xa8, 0x00, 0x66], // Ipv4(192.168.0.102)
    ///     target_hw_addr: [0xff, 0xff, 0xff, 0xff, 0xff, 0xff], // Broadcast
    ///     target_proto_addr: [0xc0, 0xa8, 0x00, 0x65], // Ipv4(192.168.0.101)
    ///     payload: [],
    /// }

    let mut ethernet_buffer: &mut [u8] = &mut [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];

    let mut ethernet_packet = MutableEthernetPacket::new(&mut ethernet_buffer).unwrap();
    ethernet_packet.set_destination(target_mac);
    ethernet_packet.set_source(source_mac);
    ethernet_packet.set_ethertype(EtherTypes::Arp);

    let mut arp_buffer: &mut [u8] = &mut [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                          0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                          0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];

    let mut arp_packet = MutableArpPacket::new(&mut arp_buffer).unwrap();

    arp_packet.set_hardware_type(ArpHardwareTypes::Ethernet);
    arp_packet.set_protocol_type(EtherTypes::Ipv4);
    arp_packet.set_hw_addr_len(6);
    arp_packet.set_proto_addr_len(4);
    arp_packet.set_operation(ArpOperations::Reply);
    arp_packet.set_sender_hw_addr(source_mac);
    arp_packet.set_sender_proto_addr(gateway);
    arp_packet.set_target_hw_addr(target_mac);
    arp_packet.set_target_proto_addr(target_ip);

    ethernet_packet.set_payload(arp_packet.packet_mut());

    tx.send_to(&ethernet_packet.to_immutable(), Some(interface));
}

fn main() {
    let config: Config = cli_main();
    let mut packet_count: i32 = 0;

    loop {
        send_arp_reply_packet(config.source_ip, config.source_mac, config.target_ip, config.target_mac);

        packet_count += 1;
        println!("Sent {} ARP reply packets.", packet_count);
        thread::sleep(Duration::new(5, 0));
    }
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate termios;

use libc::{STDIN_FILENO, STDOUT_FILENO, TIOCGWINSZ, winsize};
use termios::{Termios, tcsetattr, TCSAFLUSH, VMIN, VTIME};
use termios::{BRKINT, ICRNL, INPCK, ISTRIP, IXON};
use termios::{OPOST, CS8};
use termios::{ECHO, ICANON, IEXTEN, ISIG};

use std::cmp;
use std::env;
use std::fs::File;
use std::io::{self, Stdin, Stdout, Read, BufRead, BufReader, Error, ErrorKind, Write};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

const TAB_STOP: usize = 8;

#[derive(PartialEq, Clone, Copy)]
enum Key {
    Character(u8),
    ArrowLeft,
    ArrowRight,
    ArrowUp,
    ArrowDown,
    Delete,
    Home,
    End,
    PageUp,
    PageDown,
}

macro_rules! ctrl_key {
    ($k:expr) => ($k & 0x1f);
}

struct RawMode {
    orig_term: Termios,
}

struct Row {
    chars: String,
    render: String,
}

struct Editor {
    _mode: RawMode,
    cx: usize,
    cy: usize,
    rx: usize,
    rowoff: usize,
    coloff: usize,
    screenrows: usize,
    screencols: usize,
    rows: Option<Vec<Row>>,
    stdin: Stdin,
    stdout: Stdout,
    filename: Option<String>,
}

impl RawMode {
    fn enable_raw_mode() -> io::Result<RawMode> {
        let mut term = Termios::from_fd(STDIN_FILENO)?;
        let mode = RawMode { orig_term: term };

        term.c_iflag &= !(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
        term.c_oflag &= !OPOST;
        term.c_cflag |= CS8;
        term.c_lflag &= !(ECHO | ICANON | IEXTEN | ISIG);
        term.c_cc[VMIN] = 0;
        term.c_cc[VTIME] = 1;

        tcsetattr(STDIN_FILENO, TCSAFLUSH, &term)?;
        Ok(mode)
    }
}

impl Drop for RawMode {
    fn drop(&mut self) {
        tcsetattr(STDIN_FILENO, TCSAFLUSH, &self.orig_term).unwrap()
    }
}

impl Row {
    fn new<T>(s: T) -> Row
        where T: Into<String>
    {
        let chars = s.into();
        let mut idx = 0;
        let mut render = String::new();
        for ch in chars.chars() {
            if ch == '\t' {
                render.push(' ');
                idx += 1;
                while idx % TAB_STOP != 0 {
                    render.push(' ');
                    idx += 1;
                }
            } else {
                render.push(ch);
                idx += 1;
            }
        }
        Row { chars, render }
    }

    fn cx_to_rx(&self, cx: usize) -> usize {
        let mut rx = 0;
        for ch in self.chars.chars().take(cx) {
            if ch == '\t' {
                rx += (TAB_STOP - 1) - (rx % TAB_STOP);
            }
            rx += 1;
        }
        rx
    }
}

fn read_non_blocking<R: Read>(r: &mut R, buf: &mut [u8]) -> usize {
    r.read(buf)
        .or_else(|e| if e.kind() == ErrorKind::WouldBlock {
                     Ok(0)
                 } else {
                     Err(e)
                 })
        .expect("read_non_blocking")
}

fn byte_slice(s: &str, offset: usize, max_len: usize) -> &[u8] {
    if s.len() > max_len + offset {
        s[offset..(max_len + offset)].as_bytes()
    } else if s.len() > offset {
        s[offset..].as_bytes()
    } else {
        s[0..0].as_bytes()
    }
}

fn editor_read_key(stdin: &mut Stdin) -> Key {
    let mut buf = [0; 1];
    loop {
        let n = read_non_blocking(stdin, &mut buf);
        if n == 1 {
            if buf[0] == b'\x1b' {
                let mut seq = [0; 2];
                let n = read_non_blocking(stdin, &mut seq);
                if n == 2 && seq[0] == b'[' {
                    if seq[1] >= b'0' && seq[1] <= b'9' {
                        let mut last = [0; 1];
                        let n = read_non_blocking(stdin, &mut last);
                        if n == 1 && last[0] == b'~' {
                            return match seq[1] {
                                       b'1' | b'7' => Key::Home,
                                       b'3' => Key::Delete,
                                       b'4' | b'8' => Key::End,
                                       b'5' => Key::PageUp,
                                       b'6' => Key::PageDown,
                                       _ => Key::Character(b'\x1b'),
                                   };
                        }
                    } else {
                        return match seq[1] {
                                   b'A' => Key::ArrowUp,
                                   b'B' => Key::ArrowDown,
                                   b'C' => Key::ArrowRight,
                                   b'D' => Key::ArrowLeft,
                                   b'H' => Key::Home,
                                   b'F' => Key::End,
                                   _ => Key::Character(b'\x1b'),
                               };
                    }
                } else if n == 2 && seq[0] == b'O' {
                    return match seq[1] {
                               b'H' => Key::Home,
                               b'F' => Key::End,
                               _ => Key::Character(b'\x1b'),
                           };
                }
                return Key::Character(b'\x1b');
            } else {
                return Key::Character(buf[0]);
            }
        }
    }
}

fn get_window_size() -> io::Result<(u16, u16)> {
    let ws = winsize {
        ws_col: 0,
        ws_row: 0,
        ws_xpixel: 0,
        ws_ypixel: 0,
    };
    unsafe {
        if libc::ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == -1 || ws.ws_col == 0 {
            return Err(Error::new(ErrorKind::Other, "get_window_size: ioctl failed"));
        }
    }
    Ok((ws.ws_row, ws.ws_col))
}

impl Editor {
    fn new() -> io::Result<Editor> {
        let mode = RawMode::enable_raw_mode()?;
        let (rows, cols) = get_window_size()?;
        let stdin = io::stdin();
        let stdout = io::stdout();
        Ok(Editor {
               _mode: mode,
               cx: 0,
               cy: 0,
               rx: 0,
               rowoff: 0,
               coloff: 0,
               screenrows: (rows - 1) as usize,
               screencols: cols as usize,
               rows: None,
               stdin: stdin,
               stdout: stdout,
               filename: None,
           })
    }

    fn write(&mut self, buf: &[u8]) -> io::Result<()> {
        self.stdout.write_all(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.stdout.flush()
    }

    fn scroll(&mut self) {
        self.rx = self.cx;
        if self.cy < self.numrows() {
            self.rx = self.rows.as_ref().unwrap()[self.cy].cx_to_rx(self.cx);
        }

        if self.cy < self.rowoff {
            self.rowoff = self.cy;
        }
        if self.cy >= self.rowoff + self.screenrows {
            self.rowoff = self.cy - self.screenrows + 1;
        }
        if self.rx < self.coloff {
            self.coloff = self.rx;
        }
        if self.rx >= self.coloff + self.screencols {
            self.coloff = (1 + self.rx) - self.screencols;
        }
    }

    fn numrows(&self) -> usize {
        self.rows.as_ref().map_or(0, |v| v.len())
    }

    fn draw_rows(&mut self) -> io::Result<()> {
        let numrows = self.numrows();
        for y in 0..(self.screenrows) {
            let filerow = y + self.rowoff;
            if filerow >= numrows {
                if self.rows.is_none() && y == self.screenrows / 3 {
                    let mut msg = format!("Kilo-rs editor -- version {}", VERSION);
                    msg.truncate(self.screencols);
                    let padding = (self.screencols - msg.len()) / 2;
                    if padding > 0 {
                        self.write(b"~")?;
                        for _ in 1..padding {
                            self.write(b" ")?;
                        }
                    }
                    self.write(msg.as_bytes())?;
                } else {
                    self.write(b"~")?;
                }
            } else {
                let rows = self.rows.as_ref().unwrap();
                self.stdout
                    .write_all(byte_slice(&rows[filerow].render, self.coloff, self.screencols))?;
            }

            self.write(b"\x1b[K")?;
            self.write(b"\r\n")?;
        }
        Ok(())
    }

    fn draw_status_bar(&mut self) -> io::Result<()> {
        self.write(b"\x1b[7m")?;
        let status;
        {
            let name = self.filename
                .as_ref()
                .map_or("[No name]", |s| s.as_str());
            let mut content = format!("{:.20} - {} lines", name, self.numrows());
            content.truncate(self.screencols);
            status = content;
        }
        let rstatus = format!("{}/{}", self.cy + 1, self.numrows());

        self.write(status.as_bytes())?;
        let mut len = status.len();
        while len < self.screencols {
            if self.screencols - len == rstatus.len() {
                self.write(rstatus.as_bytes())?;
                break;
            } else {
                self.write(b" ")?;
                len += 1;
            }
        }
        self.write(b"\x1b[m")
    }

    fn refresh_screen(&mut self) -> io::Result<()> {
        self.scroll();

        self.write(b"\x1b[?25l")?;
        self.write(b"\x1b[H")?;

        self.draw_rows()?;
        self.draw_status_bar()?;

        let move_cursor = format!("\x1b[{};{}H",
                                  (self.cy - self.rowoff) + 1,
                                  (self.rx - self.coloff) + 1)
                .into_bytes();
        self.write(&move_cursor)?;

        self.write(b"\x1b[?25h")?;
        self.flush()
    }

    fn rowlen(&self, index: usize) -> usize {
        let row = self.rows.as_ref().and_then(|r| r.get(index));
        row.map_or(0, |r| r.chars.len())
    }

    fn move_cursor(&mut self, k: Key) {
        match k {
            Key::ArrowUp => {
                if self.cy > 0 {
                    self.cy -= 1;
                }
            }
            Key::ArrowDown => {
                if self.cy < self.numrows() {
                    self.cy += 1;
                }
            }
            Key::ArrowLeft => {
                if self.cx > 0 {
                    self.cx -= 1;
                } else if self.cy > 0 {
                    self.cy -= 1;
                    self.cx = self.rowlen(self.cy);
                }
            }
            Key::ArrowRight => {
                let row = self.rows.as_ref().and_then(|r| r.get(self.cy));
                let rowlen = row.map_or(0, |r| r.chars.len());
                if self.cx < rowlen {
                    self.cx += 1;
                } else if row.is_some() && self.cx == rowlen {
                    self.cx = 0;
                    self.cy += 1;
                }
            }
            _ => (),
        }

        let rowlen = self.rowlen(self.cy);
        if self.cx > rowlen {
            self.cx = rowlen;
        }
    }

    fn process_keypress(&mut self) -> bool {
        let c = editor_read_key(&mut self.stdin);

        match c {
            Key::Character(k) if k == ctrl_key!(b'q') => return false,
            Key::Home => {
                self.cx = 0;
            }
            Key::End => {
                self.cx = self.rowlen(self.cy);
            }
            Key::PageUp | Key::PageDown => {
                let key = if c == Key::PageUp {
                    self.cy = self.rowoff;
                    Key::ArrowUp
                } else {
                    self.cy = cmp::min(self.rowoff + self.screenrows - 1, self.numrows());
                    Key::ArrowDown
                };
                for _ in 0..(self.screenrows) {
                    self.move_cursor(key);
                }
            }
            Key::ArrowUp | Key::ArrowDown | Key::ArrowLeft | Key::ArrowRight => {
                self.move_cursor(c);
            }
            _ => (),
        };
        true
    }

    fn open(&mut self, filename: &str) -> io::Result<()> {
        self.filename = Some(filename.to_owned());
        let f = File::open(filename)?;
        let file = BufReader::new(&f);
        let results: io::Result<Vec<Row>> = file.lines().map(|r| r.map(Row::new)).collect();
        self.rows = Some(results?);
        Ok(())
    }
}

impl Drop for Editor {
    fn drop(&mut self) {
        // clear screen
        self.stdout.write_all(b"\x1b[2J").unwrap();
        self.stdout.write_all(b"\x1b[H").unwrap();
        self.stdout.flush().unwrap();
    }
}

fn main() {
    let mut editor = Editor::new().unwrap();
    if let Some(filename) = env::args().nth(1) {
        editor.open(&filename).unwrap();
    }

    editor.refresh_screen().unwrap();
    while editor.process_keypress() {
        editor.refresh_screen().unwrap();
    }
}
<commit_msg>Status message<commit_after>extern crate libc;
extern crate termios;

use libc::{STDIN_FILENO, STDOUT_FILENO, TIOCGWINSZ, winsize};
use termios::{Termios, tcsetattr, TCSAFLUSH, VMIN, VTIME};
use termios::{BRKINT, ICRNL, INPCK, ISTRIP, IXON};
use termios::{OPOST, CS8};
use termios::{ECHO, ICANON, IEXTEN, ISIG};

use std::borrow::Cow;
use std::cmp;
use std::env;
use std::fs::File;
use std::io::{self, Stdin, Stdout, Read, BufRead, BufReader, Error, ErrorKind, Write};
use std::time::{Duration, Instant};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

const TAB_STOP: usize = 8;

#[derive(PartialEq, Clone, Copy)]
enum Key {
    Character(u8),
    ArrowLeft,
    ArrowRight,
    ArrowUp,
    ArrowDown,
    Delete,
    Home,
    End,
    PageUp,
    PageDown,
}

macro_rules! ctrl_key {
    ($k:expr) => ($k & 0x1f);
}

struct RawMode {
    orig_term: Termios,
}

struct Row {
    chars: String,
    render: String,
}

struct Editor {
    _mode: RawMode,
    cx: usize,
    cy: usize,
    rx: usize,
    rowoff: usize,
    coloff: usize,
    screenrows: usize,
    screencols: usize,
    rows: Option<Vec<Row>>,
    stdin: Stdin,
    stdout: Stdout,
    filename: Option<String>,
    statusmsg: String,
    statusmsg_time: Instant,
}

impl RawMode {
    fn enable_raw_mode() -> io::Result<RawMode> {
        let mut term = Termios::from_fd(STDIN_FILENO)?;
        let mode = RawMode { orig_term: term };

        term.c_iflag &= !(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
        term.c_oflag &= !OPOST;
        term.c_cflag |= CS8;
        term.c_lflag &= !(ECHO | ICANON | IEXTEN | ISIG);
        term.c_cc[VMIN] = 0;
        term.c_cc[VTIME] = 1;

        tcsetattr(STDIN_FILENO, TCSAFLUSH, &term)?;
        Ok(mode)
    }
}

impl Drop for RawMode {
    fn drop(&mut self) {
        tcsetattr(STDIN_FILENO, TCSAFLUSH, &self.orig_term).unwrap()
    }
}

impl Row {
    fn new<T>(s: T) -> Row
        where T: Into<String>
    {
        let chars = s.into();
        let mut idx = 0;
        let mut render = String::new();
        for ch in chars.chars() {
            if ch == '\t' {
                render.push(' ');
                idx += 1;
                while idx % TAB_STOP != 0 {
                    render.push(' ');
                    idx += 1;
                }
            } else {
                render.push(ch);
                idx += 1;
            }
        }
        Row { chars, render }
    }

    fn cx_to_rx(&self, cx: usize) -> usize {
        let mut rx = 0;
        for ch in self.chars.chars().take(cx) {
            if ch == '\t' {
                rx += (TAB_STOP - 1) - (rx % TAB_STOP);
            }
            rx += 1;
        }
        rx
    }
}

fn read_non_blocking<R: Read>(r: &mut R, buf: &mut [u8]) -> usize {
    r.read(buf)
        .or_else(|e| if e.kind() == ErrorKind::WouldBlock {
                     Ok(0)
                 } else {
                     Err(e)
                 })
        .expect("read_non_blocking")
}

fn byte_slice(s: &str, offset: usize, max_len: usize) -> &[u8] {
    if s.len() > max_len + offset {
        s[offset..(max_len + offset)].as_bytes()
    } else if s.len() > offset {
        s[offset..].as_bytes()
    } else {
        s[0..0].as_bytes()
    }
}

fn editor_read_key(stdin: &mut Stdin) -> Key {
    let mut buf = [0; 1];
    loop {
        let n = read_non_blocking(stdin, &mut buf);
        if n == 1 {
            if buf[0] == b'\x1b' {
                let mut seq = [0; 2];
                let n = read_non_blocking(stdin, &mut seq);
                if n == 2 && seq[0] == b'[' {
                    if seq[1] >= b'0' && seq[1] <= b'9' {
                        let mut last = [0; 1];
                        let n = read_non_blocking(stdin, &mut last);
                        if n == 1 && last[0] == b'~' {
                            return match seq[1] {
                                       b'1' | b'7' => Key::Home,
                                       b'3' => Key::Delete,
                                       b'4' | b'8' => Key::End,
                                       b'5' => Key::PageUp,
                                       b'6' => Key::PageDown,
                                       _ => Key::Character(b'\x1b'),
                                   };
                        }
                    } else {
                        return match seq[1] {
                                   b'A' => Key::ArrowUp,
                                   b'B' => Key::ArrowDown,
                                   b'C' => Key::ArrowRight,
                                   b'D' => Key::ArrowLeft,
                                   b'H' => Key::Home,
                                   b'F' => Key::End,
                                   _ => Key::Character(b'\x1b'),
                               };
                    }
                } else if n == 2 && seq[0] == b'O' {
                    return match seq[1] {
                               b'H' => Key::Home,
                               b'F' => Key::End,
                               _ => Key::Character(b'\x1b'),
                           };
                }
                return Key::Character(b'\x1b');
            } else {
                return Key::Character(buf[0]);
            }
        }
    }
}

fn get_window_size() -> io::Result<(u16, u16)> {
    let ws = winsize {
        ws_col: 0,
        ws_row: 0,
        ws_xpixel: 0,
        ws_ypixel: 0,
    };
    unsafe {
        if libc::ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == -1 || ws.ws_col == 0 {
            return Err(Error::new(ErrorKind::Other, "get_window_size: ioctl failed"));
        }
    }
    Ok((ws.ws_row, ws.ws_col))
}

impl Editor {
    fn new() -> io::Result<Editor> {
        let mode = RawMode::enable_raw_mode()?;
        let (rows, cols) = get_window_size()?;
        let stdin = io::stdin();
        let stdout = io::stdout();
        Ok(Editor {
               _mode: mode,
               cx: 0,
               cy: 0,
               rx: 0,
               rowoff: 0,
               coloff: 0,
               screenrows: (rows - 2) as usize,
               screencols: cols as usize,
               rows: None,
               stdin: stdin,
               stdout: stdout,
               filename: None,
               statusmsg: String::new(),
               statusmsg_time: Instant::now(),
           })
    }

    fn write(&mut self, buf: &[u8]) -> io::Result<()> {
        self.stdout.write_all(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.stdout.flush()
    }

    fn scroll(&mut self) {
        self.rx = self.cx;
        if self.cy < self.numrows() {
            self.rx = self.rows.as_ref().unwrap()[self.cy].cx_to_rx(self.cx);
        }

        if self.cy < self.rowoff {
            self.rowoff = self.cy;
        }
        if self.cy >= self.rowoff + self.screenrows {
            self.rowoff = self.cy - self.screenrows + 1;
        }
        if self.rx < self.coloff {
            self.coloff = self.rx;
        }
        if self.rx >= self.coloff + self.screencols {
            self.coloff = (1 + self.rx) - self.screencols;
        }
    }

    fn numrows(&self) -> usize {
        self.rows.as_ref().map_or(0, |v| v.len())
    }

    fn draw_rows(&mut self) -> io::Result<()> {
        let numrows = self.numrows();
        for y in 0..(self.screenrows) {
            let filerow = y + self.rowoff;
            if filerow >= numrows {
                if self.rows.is_none() && y == self.screenrows / 3 {
                    let mut msg = format!("Kilo-rs editor -- version {}", VERSION);
                    msg.truncate(self.screencols);
                    let padding = (self.screencols - msg.len()) / 2;
                    if padding > 0 {
                        self.write(b"~")?;
                        for _ in 1..padding {
                            self.write(b" ")?;
                        }
                    }
                    self.write(msg.as_bytes())?;
                } else {
                    self.write(b"~")?;
                }
            } else {
                let rows = self.rows.as_ref().unwrap();
                self.stdout
                    .write_all(byte_slice(&rows[filerow].render, self.coloff, self.screencols))?;
            }

            self.write(b"\x1b[K")?;
            self.write(b"\r\n")?;
        }
        Ok(())
    }

    fn draw_status_bar(&mut self) -> io::Result<()> {
        self.write(b"\x1b[7m")?;
        let status;
        {
            let name = self.filename
                .as_ref()
                .map_or("[No name]", |s| s.as_str());
            let mut content = format!("{:.20} - {} lines", name, self.numrows());
            content.truncate(self.screencols);
            status = content;
        }
        let rstatus = format!("{}/{}", self.cy + 1, self.numrows());

        self.write(status.as_bytes())?;
        let mut len = status.len();
        while len < self.screencols {
            if self.screencols - len == rstatus.len() {
                self.write(rstatus.as_bytes())?;
                break;
            } else {
                self.write(b" ")?;
                len += 1;
            }
        }
        self.write(b"\x1b[m")?;
        self.write(b"\r\n")
    }

    fn draw_message_bar(&mut self) -> io::Result<()> {
        self.write(b"\x1b[K")?;
        if !self.statusmsg.is_empty() && self.statusmsg_time.elapsed() < Duration::from_secs(5) {
            let mut msg = Cow::from(self.statusmsg.as_str());
            if msg.len() > self.screencols {
                msg.to_mut().truncate(self.screencols);
            }
            self.stdout.write_all(msg.as_bytes())?;
        }
        Ok(())
    }

    fn refresh_screen(&mut self) -> io::Result<()> {
        self.scroll();

        self.write(b"\x1b[?25l")?;
        self.write(b"\x1b[H")?;

        self.draw_rows()?;
        self.draw_status_bar()?;
        self.draw_message_bar()?;

        let move_cursor = format!("\x1b[{};{}H",
                                  (self.cy - self.rowoff) + 1,
                                  (self.rx - self.coloff) + 1)
                .into_bytes();
        self.write(&move_cursor)?;

        self.write(b"\x1b[?25h")?;
        self.flush()
    }

    fn set_status_message<S: Into<String>>(&mut self, msg: S) {
        self.statusmsg = msg.into();
        self.statusmsg_time = Instant::now();
    }

    fn rowlen(&self, index: usize) -> usize {
        let row = self.rows.as_ref().and_then(|r| r.get(index));
        row.map_or(0, |r| r.chars.len())
    }

    fn move_cursor(&mut self, k: Key) {
        match k {
            Key::ArrowUp => {
                if self.cy > 0 {
                    self.cy -= 1;
                }
            }
            Key::ArrowDown => {
                if self.cy < self.numrows() {
                    self.cy += 1;
                }
            }
            Key::ArrowLeft => {
                if self.cx > 0 {
                    self.cx -= 1;
                } else if self.cy > 0 {
                    self.cy -= 1;
                    self.cx = self.rowlen(self.cy);
                }
            }
            Key::ArrowRight => {
                let row = self.rows.as_ref().and_then(|r| r.get(self.cy));
                let rowlen = row.map_or(0, |r| r.chars.len());
                if self.cx < rowlen {
                    self.cx += 1;
                } else if row.is_some() && self.cx == rowlen {
                    self.cx = 0;
                    self.cy += 1;
                }
            }
            _ => (),
        }

        let rowlen = self.rowlen(self.cy);
        if self.cx > rowlen {
            self.cx = rowlen;
        }
    }

    fn process_keypress(&mut self) -> bool {
        let c = editor_read_key(&mut self.stdin);

        match c {
            Key::Character(k) if k == ctrl_key!(b'q') => return false,
            Key::Home => {
                self.cx = 0;
            }
            Key::End => {
                self.cx = self.rowlen(self.cy);
            }
            Key::PageUp | Key::PageDown => {
                let key = if c == Key::PageUp {
                    self.cy = self.rowoff;
                    Key::ArrowUp
                } else {
                    self.cy = cmp::min(self.rowoff + self.screenrows - 1, self.numrows());
                    Key::ArrowDown
                };
                for _ in 0..(self.screenrows) {
                    self.move_cursor(key);
                }
            }
            Key::ArrowUp | Key::ArrowDown | Key::ArrowLeft | Key::ArrowRight => {
                self.move_cursor(c);
            }
            _ => (),
        };
        true
    }

    fn open(&mut self, filename: &str) -> io::Result<()> {
        self.filename = Some(filename.to_owned());
        let f = File::open(filename)?;
        let file = BufReader::new(&f);
        let results: io::Result<Vec<Row>> = file.lines().map(|r| r.map(Row::new)).collect();
        self.rows = Some(results?);
        Ok(())
    }
}

impl Drop for Editor {
    fn drop(&mut self) {
        // clear screen
        self.stdout.write_all(b"\x1b[2J").unwrap();
        self.stdout.write_all(b"\x1b[H").unwrap();
        self.stdout.flush().unwrap();
    }
}

fn main() {
    let mut editor = Editor::new().unwrap();
    if let Some(filename) = env::args().nth(1) {
        editor.open(&filename).unwrap();
    }

    editor.set_status_message("HELP: Ctrl-Q = quit");

    editor.refresh_screen().unwrap();
    while editor.process_keypress() {
        editor.refresh_screen().unwrap();
    }
}
<|endoftext|>


<commit_before>use clap::Parser;
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(":", "")
        .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.byte_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(b"");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(
                row.iter().map(String::from_utf8_lossy),
            ))
            .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 col_count = stmt.column_count();

    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..col_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 col_count = stmt.column_count();

    let mut writer = csv::Writer::from_path(destination_path).map_err(|e| format!("{:?}", e))?;
    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..col_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)
        },
    )?;
    c.create_scalar_function(
        "regexp",
        2,
        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)?;
            Ok(re.is_match(&value))
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 5] = [b' ', 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::Parser)]
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 = [
        // keywords
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        // functions
        "length",
        "coalesce",
        "regexp_extract",
        // csv-sql commands
        "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));
    let history_path = dirs::home_dir().unwrap().join(".csv-sql-history");
    let _ = rl.load_history(&history_path);
    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;
            }
        }
    }
    rl.save_history(&history_path).unwrap();

    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>satisfy clippy (#670)<commit_after>use clap::Parser;
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(':', "")
        .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.byte_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(b"");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(
                row.iter().map(String::from_utf8_lossy),
            ))
            .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 col_count = stmt.column_count();

    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..col_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 col_count = stmt.column_count();

    let mut writer = csv::Writer::from_path(destination_path).map_err(|e| format!("{:?}", e))?;
    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..col_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)
        },
    )?;
    c.create_scalar_function(
        "regexp",
        2,
        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)?;
            Ok(re.is_match(&value))
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 5] = [b' ', 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::Parser)]
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 = [
        // keywords
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        // functions
        "length",
        "coalesce",
        "regexp_extract",
        // csv-sql commands
        "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));
    let history_path = dirs::home_dir().unwrap().join(".csv-sql-history");
    let _ = rl.load_history(&history_path);
    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;
            }
        }
    }
    rl.save_history(&history_path).unwrap();

    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|>