<commit_before>#[macro_use(crate_version, crate_authors)]
extern crate clap;

extern crate api_server;
extern crate devices;
extern crate sys_util;
extern crate vmm;

use clap::{App, Arg};
use std::fs::File;
use std::path::PathBuf;
use std::sync::mpsc::channel;

use api_server::ApiServer;
use sys_util::{syslog, EventFd, GuestAddress};
use vmm::{CMDLINE_MAX_SIZE, CMDLINE_OFFSET, KERNEL_START_OFFSET};
use vmm::{KernelConfig, kernel_cmdline};
use vmm::machine::MachineCfg;
use vmm::device_config::BlockDeviceConfig;

fn main() {
    if let Err(e) = syslog::init() {
        println!("failed to initialize syslog: {:?}", e);
        return;
    }

    let cmd_arguments = App::new("firecracker")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Launch a microvm.")
        .arg(
            Arg::with_name("vmm_no_api")
                .long("vmm-no-api")
                .help("Start vmm by default, no API thread")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("api_sock")
                .long("api-sock")
                .help("Path to unix domain socket used by the API")
                .default_value("/tmp/firecracker.socket")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("kernel_path")
                .short("k")
                .long("kernel-path")
                .help("The kernel's file path (vmlinux.bin)")
                .required(true)
                .takes_value(true),
        )
        .arg(
            Arg::with_name("kernel_cmdline")
                .long("kernel-cmdline")
                .help("The kernel's command line")
                .default_value("console=ttyS0 noapic reboot=k panic=1 pci=off nomodules")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("mem_size")
                .long("mem-size")
                .default_value("128")
                .help("Virtual Machine Memory Size in MiB")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("vcpu_count")
                .long("vcpu-count")
                .default_value("1")
                .help("Number of VCPUs")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("root_blk_file")
                .short("r")
                .long("root-blk")
                .help("File to serve as root block device [enabled only with --vmm-no-api]")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("host_ip")
                .long("host-ip")
                .help("IPv4 address of the host interface")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("subnet_mask")
                .long("subnet-mask")
                .default_value("255.255.255.0")
                .help("Subnet mask for the IP address of host interface")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("vsock_guest_cid")
                .long("vsock-guest-cid")
                .help("The guest CID for the virtio-vhost-vsock device")
                .takes_value(true)
        )
        /*
        The mac address option is not currently implemented; the L2 addresses for both the
        host interface and the guest interface take some implicit (possibly random) values
        .arg(
            Arg::with_name("mac_address")
                .long("mac-addr")
                .help("MAC address for the VM")
                .takes_value(true)
        )*/
        .get_matches();

    let vmm_no_api = cmd_arguments.is_present("vmm_no_api");

    let cfg = parse_args(&cmd_arguments);
    let (to_vmm, from_api) = channel();

    // TODO: vmm_no_api is for integration testing, need to find a more pretty solution
    if vmm_no_api {
        let mut vmm = vmm::Vmm::new(
            cfg,
            EventFd::new().expect("cannot create eventFD"),
            from_api,
        ).expect("cannot create VMM");
        // This is a temporary fix. Block devices should be added via http requests.
        // With the command line, we can only add one device, with default to root block device.
        if cmd_arguments.is_present("root_blk_file") {
            let root_block_device = BlockDeviceConfig {
                path_on_host: PathBuf::from(cmd_arguments.value_of("root_blk_file").unwrap()),
                is_root_device: true,
                drive_id: String::from("1"),
            };
            vmm.put_block_device(root_block_device).expect("cannot add root block device.");
        }
        // configure kernel from command line
        //we're using unwrap here because the kernel_path is mandatory for now
        let kernel_file = File::open(cmd_arguments.value_of("kernel_path")
            .unwrap_or_default())
            .expect("Cannot open kernel file");

        let mut cmdline = kernel_cmdline::Cmdline::new(CMDLINE_MAX_SIZE);
        cmdline
            .insert_str(cmd_arguments.value_of("kernel_cmdline").unwrap())
            .expect("could not use the specified kernel cmdline");
        let kernel_config = KernelConfig {
            cmdline,
            kernel_file,
            kernel_start_addr: GuestAddress(KERNEL_START_OFFSET),
            cmdline_addr: GuestAddress(CMDLINE_OFFSET)
        };
        vmm.configure_kernel(kernel_config);
        vmm.boot_kernel().expect("cannot boot kernel");
        vmm.run_control().expect("VMM loop error!");
    } else {
        // safe to unwrap since api_sock has a default value
        let bind_path = cmd_arguments
            .value_of("api_sock")
            .map(|s| PathBuf::from(s))
            .unwrap();

        let server = ApiServer::new(to_vmm, 100).unwrap();
        let api_event_fd = server
            .get_event_fd_clone()
            .expect("cannot clone API eventFD");
        let _vmm_thread_handle = vmm::start_vmm_thread(cfg, api_event_fd, from_api);
        server.bind_and_run(bind_path).unwrap();
    }
}

fn parse_args(cmd_arguments: &clap::ArgMatches) -> MachineCfg {

    let vcpu_count = match cmd_arguments
        .value_of("vcpu_count")
        .unwrap()
        .to_string()
        .parse::<u8>()
    {
        Ok(value) => value,
        Err(error) => {
            panic!("Invalid value for vcpu_count! {:?}", error);
        }
    };

    let mem_size = match cmd_arguments
        .value_of("mem_size")
        .unwrap()
        .to_string()
        .parse::<usize>()
    {
        Ok(value) => value,
        Err(error) => {
            panic!("Invalid value for mem_size! {:?}", error);
        }
    };

    let root_blk_file = cmd_arguments
        .value_of("root_blk_file")
        .map(|s| PathBuf::from(s));

    //fixme print some message when the Ipv4Addrs cannot be parsed
    let host_ip = cmd_arguments
        .value_of("host_ip")
        .map(|x| x.parse().unwrap());

    let subnet_mask = cmd_arguments
        .value_of("subnet_mask")
        .unwrap()
        .parse()
        .unwrap();

    let vsock_guest_cid = match cmd_arguments.value_of("vsock_guest_cid") {
        Some(cid) => match cid.parse::<u64>() {
            Ok(value) => Some(value),
            Err(error) => {
                panic!(
                    "Invalid parameter value for the vsock's guest CID! {:?}",
                    error
                );
            }
        },
        None => None,
    };

    MachineCfg::new(
        vcpu_count,
        mem_size,
        root_blk_file,
        host_ip,
        subnet_mask,
        vsock_guest_cid,
    )
}
<commit_msg>main: refactoring main function for vmm_no_api<commit_after>#[macro_use(crate_version, crate_authors)]
extern crate clap;

extern crate api_server;
extern crate devices;
extern crate sys_util;
extern crate vmm;

use clap::{App, Arg};
use std::fs::File;
use std::path::PathBuf;
use std::sync::mpsc::channel;
use std::sync::mpsc::Receiver;

use api_server::{ApiRequest, ApiServer};
use sys_util::{syslog, EventFd, GuestAddress};
use vmm::{CMDLINE_MAX_SIZE, CMDLINE_OFFSET, KERNEL_START_OFFSET};
use vmm::{kernel_cmdline, KernelConfig};
use vmm::machine::MachineCfg;
use vmm::device_config::BlockDeviceConfig;

fn main() {
    if let Err(e) = syslog::init() {
        println!("failed to initialize syslog: {:?}", e);
        return;
    }

    let cmd_arguments = App::new("firecracker")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Launch a microvm.")
        .arg(
            Arg::with_name("vmm_no_api")
                .long("vmm-no-api")
                .help("Start vmm by default, no API thread")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("api_sock")
                .long("api-sock")
                .help("Path to unix domain socket used by the API")
                .default_value("/tmp/firecracker.socket")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("kernel_path")
                .short("k")
                .long("kernel-path")
                .help("The kernel's file path (vmlinux.bin)")
                .required(true)
                .takes_value(true),
        )
        .arg(
            Arg::with_name("kernel_cmdline")
                .long("kernel-cmdline")
                .help("The kernel's command line")
                .default_value("console=ttyS0 noapic reboot=k panic=1 pci=off nomodules")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("mem_size")
                .long("mem-size")
                .default_value("128")
                .help("Virtual Machine Memory Size in MiB")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("vcpu_count")
                .long("vcpu-count")
                .default_value("1")
                .help("Number of VCPUs")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("root_blk_file")
                .short("r")
                .long("root-blk")
                .help("File to serve as root block device [enabled only with --vmm-no-api]")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("host_ip")
                .long("host-ip")
                .help("IPv4 address of the host interface")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("subnet_mask")
                .long("subnet-mask")
                .default_value("255.255.255.0")
                .help("Subnet mask for the IP address of host interface")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("vsock_guest_cid")
                .long("vsock-guest-cid")
                .help("The guest CID for the virtio-vhost-vsock device")
                .takes_value(true)
        )
        /*
        The mac address option is not currently implemented; the L2 addresses for both the
        host interface and the guest interface take some implicit (possibly random) values
        .arg(
            Arg::with_name("mac_address")
                .long("mac-addr")
                .help("MAC address for the VM")
                .takes_value(true)
        )*/
        .get_matches();

    let vmm_no_api = cmd_arguments.is_present("vmm_no_api");

    let cfg = parse_args(&cmd_arguments);
    let (to_vmm, from_api) = channel();

    // TODO: vmm_no_api is for integration testing, need to find a more pretty solution
    if vmm_no_api {
        vmm_no_api_handler(cmd_arguments, cfg, from_api);
    } else {
        // safe to unwrap since api_sock has a default value
        let bind_path = cmd_arguments
            .value_of("api_sock")
            .map(|s| PathBuf::from(s))
            .unwrap();

        let server = ApiServer::new(to_vmm, 100).unwrap();
        let api_event_fd = server
            .get_event_fd_clone()
            .expect("cannot clone API eventFD");
        let _vmm_thread_handle = vmm::start_vmm_thread(cfg, api_event_fd, from_api);
        server.bind_and_run(bind_path).unwrap();
    }
}

fn vmm_no_api_handler(
    cmd_arguments: clap::ArgMatches,
    cfg: MachineCfg,
    from_api: Receiver<Box<ApiRequest>>,
) {
    let mut vmm = vmm::Vmm::new(
        cfg,
        EventFd::new().expect("cannot create eventFD"),
        from_api,
    ).expect("cannot create VMM");
    // This is a temporary fix. Block devices should be added via http requests.
    // With the command line, we can only add one device, with default to root block device.
    if cmd_arguments.is_present("root_blk_file") {
        let root_block_device = BlockDeviceConfig {
            path_on_host: PathBuf::from(cmd_arguments.value_of("root_blk_file").unwrap()),
            is_root_device: true,
            drive_id: String::from("1"),
        };
        vmm.put_block_device(root_block_device)
            .expect("cannot add root block device.");
    }
    // configure kernel from command line
    //we're using unwrap here because the kernel_path is mandatory for now
    let kernel_file = File::open(cmd_arguments.value_of("kernel_path").unwrap_or_default())
        .expect("Cannot open kernel file");

    let mut cmdline = kernel_cmdline::Cmdline::new(CMDLINE_MAX_SIZE);
    cmdline
        .insert_str(cmd_arguments.value_of("kernel_cmdline").unwrap())
        .expect("could not use the specified kernel cmdline");
    let kernel_config = KernelConfig {
        cmdline,
        kernel_file,
        kernel_start_addr: GuestAddress(KERNEL_START_OFFSET),
        cmdline_addr: GuestAddress(CMDLINE_OFFSET),
    };
    vmm.configure_kernel(kernel_config);
    vmm.boot_kernel().expect("cannot boot kernel");
    vmm.run_control().expect("VMM loop error!");
}

fn parse_args(cmd_arguments: &clap::ArgMatches) -> MachineCfg {
    let vcpu_count = match cmd_arguments
        .value_of("vcpu_count")
        .unwrap()
        .to_string()
        .parse::<u8>()
    {
        Ok(value) => value,
        Err(error) => {
            panic!("Invalid value for vcpu_count! {:?}", error);
        }
    };

    let mem_size = match cmd_arguments
        .value_of("mem_size")
        .unwrap()
        .to_string()
        .parse::<usize>()
    {
        Ok(value) => value,
        Err(error) => {
            panic!("Invalid value for mem_size! {:?}", error);
        }
    };

    let root_blk_file = cmd_arguments
        .value_of("root_blk_file")
        .map(|s| PathBuf::from(s));

    //fixme print some message when the Ipv4Addrs cannot be parsed
    let host_ip = cmd_arguments
        .value_of("host_ip")
        .map(|x| x.parse().unwrap());

    let subnet_mask = cmd_arguments
        .value_of("subnet_mask")
        .unwrap()
        .parse()
        .unwrap();

    let vsock_guest_cid = match cmd_arguments.value_of("vsock_guest_cid") {
        Some(cid) => match cid.parse::<u64>() {
            Ok(value) => Some(value),
            Err(error) => {
                panic!(
                    "Invalid parameter value for the vsock's guest CID! {:?}",
                    error
                );
            }
        },
        None => None,
    };

    MachineCfg::new(
        vcpu_count,
        mem_size,
        root_blk_file,
        host_ip,
        subnet_mask,
        vsock_guest_cid,
    )
}
<|endoftext|>


<commit_before>#![feature(globs)]

extern crate libc;
extern crate neovim;
extern crate rgtk;
extern crate serialize;

use neovim::nvim_main;
use rgtk::*;
use std::collections::HashSet;
use std::{io, str, time};

mod projects;
mod ui;
mod utils;

mod ffi {
    pub use libc::{c_char, c_int, c_uchar, c_void, uint64_t};
    pub use libc::funcs::posix88::unistd::{close, pipe, read, write};
    pub use libc::types::os::arch::c95::size_t;

    extern "C" {
        pub fn fork () -> c_int;
        pub fn kill (pid: c_int, sig: c_int);
        pub fn channel_from_fds (read_fd: c_int, write_fd: c_int) -> uint64_t;
        pub fn channel_subscribe (id: uint64_t, event: *const c_char);
    }
}

fn gui_main(
    pty: &mut gtk::VtePty,
    read_fd: ffi::c_int,
    write_fd: ffi::c_int,
    pid: ffi::c_int)
{
    gtk::init();

    // constants

    let width = 1242;
    let height = 768;
    let editor_height = ((height as f32) * 0.75) as i32;

    // create the window

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    window.set_title("SolidOak");
    window.set_window_position(gtk::WindowPosition::Center);
    window.set_default_size(width, height);

    window.connect(gtk::signals::DeleteEvent::new(|_| {
        unsafe {
            ffi::close(read_fd);
            ffi::close(write_fd);
            ffi::kill(pid, 15);
        }
        gtk::main_quit();
        true
    }));

    // create the panes

    let new_button = gtk::Button::new_with_label("New Project").unwrap();
    let import_button = gtk::Button::new_with_label("Import").unwrap();
    let rename_button = gtk::Button::new_with_label("Rename").unwrap();
    let remove_button = gtk::Button::new_with_label("Remove").unwrap();

    let mut proj_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    proj_btns.set_size_request(-1, -1);
    proj_btns.add(&new_button);
    proj_btns.add(&import_button);
    proj_btns.add(&rename_button);
    proj_btns.add(&remove_button);

    let mut proj_tree = gtk::TreeView::new().unwrap();
    let selection = proj_tree.get_selection().unwrap();
    let column_types = [glib::ffi::g_type_string, glib::ffi::g_type_string];
    let store = gtk::TreeStore::new(&column_types).unwrap();
    let model = store.get_model().unwrap();
    proj_tree.set_model(&model);
    proj_tree.set_headers_visible(false);

    let mut scroll_pane = gtk::ScrolledWindow::new(None, None).unwrap();
    scroll_pane.add(&proj_tree);

    let column = gtk::TreeViewColumn::new().unwrap();
    let cell = gtk::CellRendererText::new().unwrap();
    column.pack_start(&cell, true);
    column.add_attribute(&cell, "text", 0);
    proj_tree.append_column(&column);

    let mut proj_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    proj_pane.set_size_request(-1, -1);
    proj_pane.pack_start(&proj_btns, false, true, 0);
    proj_pane.pack_start(&scroll_pane, true, true, 0);

    let mut editor_pane = gtk::VteTerminal::new().unwrap();
    editor_pane.set_size_request(-1, editor_height);
    editor_pane.set_pty(pty);
    editor_pane.watch_child(pid);

    let run_button = gtk::Button::new_with_label("Run").unwrap();
    let build_button = gtk::Button::new_with_label("Build").unwrap();
    let test_button = gtk::Button::new_with_label("Test").unwrap();
    let clean_button = gtk::Button::new_with_label("Clean").unwrap();
    let stop_button = gtk::Button::new_with_label("Stop").unwrap();

    let mut build_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    build_btns.set_size_request(-1, -1);
    build_btns.add(&run_button);
    build_btns.add(&build_button);
    build_btns.add(&test_button);
    build_btns.add(&clean_button);
    build_btns.add(&stop_button);

    let build_term = gtk::VteTerminal::new().unwrap();

    let mut build_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    build_pane.pack_start(&build_btns, false, true, 0);
    build_pane.pack_start(&build_term, true, true, 0);

    let mut content = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    content.pack_start(&editor_pane, false, true, 0);
    content.pack_start(&build_pane, true, true, 0);

    let mut hbox = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    hbox.pack_start(&proj_pane, false, true, 0);
    hbox.pack_start(&content, true, true, 0);
    window.add(&hbox);

    // populate the project tree

    let mut state = ::utils::State{
        projects: HashSet::new(),
        expansions: HashSet::new(),
        selection: None,
        tree_model: &model,
        tree_store: &store,
        tree_selection: &selection,
        rename_button: &rename_button,
        remove_button: &remove_button,
    };

    ::utils::create_data_dir();
    ::utils::read_prefs(&mut state);
    ::ui::update_project_tree(&mut state, &mut proj_tree);

    // connect to the signals

    new_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::new_project(&mut state, &mut proj_tree);
    }));
    import_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::import_project(&mut state, &mut proj_tree);
    }));
    rename_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::rename_file(&mut state);
    }));
    remove_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::remove_item(&mut state);
    }));
    selection.connect(gtk::signals::Changed::new(|| {
        ::projects::update_selection(&mut state);
    }));
    proj_tree.connect(gtk::signals::RowCollapsed::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::remove_expansion(&mut state, &iter);
    }));
    proj_tree.connect(gtk::signals::RowExpanded::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::add_expansion(&mut state, &iter);
    }));

    // show the window

    window.show_all();
    gtk::main();
}

fn main() {
    // takes care of piping stdin/stdout between the gui and nvim
    let mut pty = gtk::VtePty::new().unwrap();

    // two pairs of anonymous pipes for msgpack-rpc between the gui and nvim
    let mut nvim_from_gui : [ffi::c_int, ..2] = [0, ..2];
    let mut gui_from_nvim : [ffi::c_int, ..2] = [0, ..2];
    unsafe {
        ffi::pipe(nvim_from_gui.as_mut_ptr());
        ffi::pipe(gui_from_nvim.as_mut_ptr());
    };

    // split into two processes
    let pid = unsafe { ffi::fork() };

    if pid > 0 { // the gui process
        // listen for messages from nvim
        spawn(proc() {
            let mut buf : [ffi::c_uchar, ..100] = [0, ..100];
            unsafe {
                loop {
                    let buf_ptr = buf.as_mut_ptr() as *mut ffi::c_void;
                    let n = ffi::read(gui_from_nvim[0], buf_ptr, 100);
                    if n < 0 { break; }
                    let msg = str::from_utf8(buf.slice_to(n as uint)).unwrap();
                    println!("Received: {}", msg);
                }
            }
        });

        // start the gui
        gui_main(&mut pty, gui_from_nvim[0], gui_from_nvim[1], pid);
    } else { // the nvim process
        // prepare this process to be piped into the gui
        pty.child_setup();

        // listen for messages from the gui
        spawn(proc() {
            io::timer::sleep(time::Duration::seconds(1));
            unsafe {
                let id = ffi::channel_from_fds(nvim_from_gui[0], gui_from_nvim[1]);
                let cmd = "test";
                ffi::channel_subscribe(id, cmd.to_c_str().as_ptr());

                let msg = "Hello, world!";
                let msg_c = msg.to_c_str();
                let msg_ptr = msg_c.as_ptr() as *const ffi::c_void;
                ffi::write(gui_from_nvim[1], msg_ptr, msg_c.len() as ffi::size_t);
            }
        });

        // start nvim
        nvim_main();
    }
}
<commit_msg>Move neovim ffi out of project<commit_after>#![feature(globs)]

extern crate libc;
extern crate neovim;
extern crate rgtk;
extern crate serialize;

use neovim::nvim_main;
use rgtk::*;
use std::collections::HashSet;
use std::{io, str, time};

mod projects;
mod ui;
mod utils;

mod ffi {
    pub use libc::{c_int, c_uchar, c_void};
    pub use libc::funcs::posix88::unistd::{close, pipe, read, write};
    pub use libc::types::os::arch::c95::size_t;

    extern "C" {
        pub fn fork () -> c_int;
        pub fn kill (pid: c_int, sig: c_int);
    }
}

fn gui_main(
    pty: &mut gtk::VtePty,
    read_fd: ffi::c_int,
    write_fd: ffi::c_int,
    pid: ffi::c_int)
{
    gtk::init();

    // constants

    let width = 1242;
    let height = 768;
    let editor_height = ((height as f32) * 0.75) as i32;

    // create the window

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    window.set_title("SolidOak");
    window.set_window_position(gtk::WindowPosition::Center);
    window.set_default_size(width, height);

    window.connect(gtk::signals::DeleteEvent::new(|_| {
        unsafe {
            ffi::close(read_fd);
            ffi::close(write_fd);
            ffi::kill(pid, 15);
        }
        gtk::main_quit();
        true
    }));

    // create the panes

    let new_button = gtk::Button::new_with_label("New Project").unwrap();
    let import_button = gtk::Button::new_with_label("Import").unwrap();
    let rename_button = gtk::Button::new_with_label("Rename").unwrap();
    let remove_button = gtk::Button::new_with_label("Remove").unwrap();

    let mut proj_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    proj_btns.set_size_request(-1, -1);
    proj_btns.add(&new_button);
    proj_btns.add(&import_button);
    proj_btns.add(&rename_button);
    proj_btns.add(&remove_button);

    let mut proj_tree = gtk::TreeView::new().unwrap();
    let selection = proj_tree.get_selection().unwrap();
    let column_types = [glib::ffi::g_type_string, glib::ffi::g_type_string];
    let store = gtk::TreeStore::new(&column_types).unwrap();
    let model = store.get_model().unwrap();
    proj_tree.set_model(&model);
    proj_tree.set_headers_visible(false);

    let mut scroll_pane = gtk::ScrolledWindow::new(None, None).unwrap();
    scroll_pane.add(&proj_tree);

    let column = gtk::TreeViewColumn::new().unwrap();
    let cell = gtk::CellRendererText::new().unwrap();
    column.pack_start(&cell, true);
    column.add_attribute(&cell, "text", 0);
    proj_tree.append_column(&column);

    let mut proj_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    proj_pane.set_size_request(-1, -1);
    proj_pane.pack_start(&proj_btns, false, true, 0);
    proj_pane.pack_start(&scroll_pane, true, true, 0);

    let mut editor_pane = gtk::VteTerminal::new().unwrap();
    editor_pane.set_size_request(-1, editor_height);
    editor_pane.set_pty(pty);
    editor_pane.watch_child(pid);

    let run_button = gtk::Button::new_with_label("Run").unwrap();
    let build_button = gtk::Button::new_with_label("Build").unwrap();
    let test_button = gtk::Button::new_with_label("Test").unwrap();
    let clean_button = gtk::Button::new_with_label("Clean").unwrap();
    let stop_button = gtk::Button::new_with_label("Stop").unwrap();

    let mut build_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    build_btns.set_size_request(-1, -1);
    build_btns.add(&run_button);
    build_btns.add(&build_button);
    build_btns.add(&test_button);
    build_btns.add(&clean_button);
    build_btns.add(&stop_button);

    let build_term = gtk::VteTerminal::new().unwrap();

    let mut build_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    build_pane.pack_start(&build_btns, false, true, 0);
    build_pane.pack_start(&build_term, true, true, 0);

    let mut content = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    content.pack_start(&editor_pane, false, true, 0);
    content.pack_start(&build_pane, true, true, 0);

    let mut hbox = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    hbox.pack_start(&proj_pane, false, true, 0);
    hbox.pack_start(&content, true, true, 0);
    window.add(&hbox);

    // populate the project tree

    let mut state = ::utils::State{
        projects: HashSet::new(),
        expansions: HashSet::new(),
        selection: None,
        tree_model: &model,
        tree_store: &store,
        tree_selection: &selection,
        rename_button: &rename_button,
        remove_button: &remove_button,
    };

    ::utils::create_data_dir();
    ::utils::read_prefs(&mut state);
    ::ui::update_project_tree(&mut state, &mut proj_tree);

    // connect to the signals

    new_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::new_project(&mut state, &mut proj_tree);
    }));
    import_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::import_project(&mut state, &mut proj_tree);
    }));
    rename_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::rename_file(&mut state);
    }));
    remove_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::remove_item(&mut state);
    }));
    selection.connect(gtk::signals::Changed::new(|| {
        ::projects::update_selection(&mut state);
    }));
    proj_tree.connect(gtk::signals::RowCollapsed::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::remove_expansion(&mut state, &iter);
    }));
    proj_tree.connect(gtk::signals::RowExpanded::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::add_expansion(&mut state, &iter);
    }));

    // show the window

    window.show_all();
    gtk::main();
}

fn main() {
    // takes care of piping stdin/stdout between the gui and nvim
    let mut pty = gtk::VtePty::new().unwrap();

    // two pairs of anonymous pipes for msgpack-rpc between the gui and nvim
    let mut nvim_from_gui : [ffi::c_int, ..2] = [0, ..2];
    let mut gui_from_nvim : [ffi::c_int, ..2] = [0, ..2];
    unsafe {
        ffi::pipe(nvim_from_gui.as_mut_ptr());
        ffi::pipe(gui_from_nvim.as_mut_ptr());
    };

    // split into two processes
    let pid = unsafe { ffi::fork() };

    if pid > 0 { // the gui process
        // listen for messages from nvim
        spawn(proc() {
            let mut buf : [ffi::c_uchar, ..100] = [0, ..100];
            unsafe {
                loop {
                    let buf_ptr = buf.as_mut_ptr() as *mut ffi::c_void;
                    let n = ffi::read(gui_from_nvim[0], buf_ptr, 100);
                    if n < 0 { break; }
                    let msg = str::from_utf8(buf.slice_to(n as uint)).unwrap();
                    println!("Received: {}", msg);
                }
            }
        });

        // start the gui
        gui_main(&mut pty, gui_from_nvim[0], gui_from_nvim[1], pid);
    } else { // the nvim process
        // prepare this process to be piped into the gui
        pty.child_setup();

        // listen for messages from the gui
        spawn(proc() {
            io::timer::sleep(time::Duration::seconds(1));

            let mut ch = neovim::Channel::new(nvim_from_gui[0], gui_from_nvim[1]);
            ch.subscribe("test");

            unsafe {
                let msg = "Hello, world!";
                let msg_c = msg.to_c_str();
                let msg_ptr = msg_c.as_ptr() as *const ffi::c_void;
                ffi::write(gui_from_nvim[1], msg_ptr, msg_c.len() as ffi::size_t);
            }
        });

        // start nvim
        nvim_main();
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate git2;

use git2::{Error, ObjectType, Repository, Commit, Oid};
use clap::App;
use std::collections::HashMap;

fn main() {
    let yaml = load_yaml!("cli.yaml");
    let matches = App::from_yaml(yaml).get_matches();

    let path = matches.value_of("path").expect("Could not parse 'path' parameter");
    let revision_range = matches.value_of("revision_range").expect("Could not parse 'revision range' parameter");

    let repo = Repository::open(path).expect("Could not open repository");
    let tags = get_tags(&repo).expect("Could not retrieve tags from repo");
    parse_log(&repo, revision_range, &tags).expect("Could not parse log");
}

type TagHashMap = HashMap<Oid, String>;
fn get_tags(repo: &Repository) -> Result<TagHashMap, Error> {
    let mut tags = HashMap::new();
    for name in try!(repo.tag_names(None)).iter() {
        let name = name.expect("Could not retrieve tag name");
        let obj = try!(repo.revparse_single(name));
        if let Ok(tag) = obj.into_tag() {
            let tag_name = tag.name().expect("Could not parse tag name").to_owned();
            tags.insert(tag.target_id(), tag_name);
        }
    }
    Ok(tags)
}

fn parse_log(repo: &Repository, revision_range: &str, tags: &TagHashMap) -> Result<(), Error> {
    let mut revwalk = try!(repo.revwalk());
    revwalk.set_sorting(git2::SORT_TIME);

    // Fill the revwalk with the selected revisions.
    let revspec = try!(repo.revparse(&revision_range));
    if revspec.mode().contains(git2::REVPARSE_SINGLE) {
        try!(revwalk.push(revspec.from().unwrap().id()));
    } else {
        let from = revspec.from().unwrap().id();
        let to = revspec.to().unwrap().id();
        try!(revwalk.push(to));
        if revspec.mode().contains(git2::REVPARSE_MERGE_BASE) {
            let base = try!(repo.merge_base(from, to));
            let o = try!(repo.find_object(base, Some(ObjectType::Commit)));
            try!(revwalk.push(o.id()));
        }
        try!(revwalk.hide(from));
    }

    // Iterate over the git objects and process them.
    for id in revwalk {
        let oid = try!(id);
        let commit = try!(repo.find_commit(oid));
        if let Some(tag) = tags.get(&oid) {
            println!("Found tag {} for commit {}", tag, oid);
        }
        print_commit(commit);
    }
    Ok(())
}

fn print_commit(commit: Commit) {
    println!("Commit: {}", commit.id());

    if commit.parents().len() > 1 {
        print!("Merge:");
        for id in commit.parent_ids() {
            print!(" {:.8}", id);
        }
        println!("");
    }

    let author = commit.author();
    println!("Author: {}\n", author);

    for line in String::from_utf8_lossy(commit.message_bytes()).lines() {
        println!("\t{}", line);
    }
    println!("\n");
}
<commit_msg>Changed structure to use a dedicated GitJournal object<commit_after>#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate git2;

use git2::{Error, ObjectType, Repository, Commit, Oid};
use clap::App;

struct GitJournal {
    repo: Repository,
    tags: Vec<(Oid, String)>,
}

impl GitJournal {
    fn new(path: &str) -> Result<GitJournal, Error> {
        // Open the repository
        let new_repo = try!(Repository::open(path));

        // Get all available tags
        let mut new_tags = vec![];
        for name in try!(new_repo.tag_names(None)).iter() {
            let name = try!(name.ok_or(Error::from_str("Could not receive tag name")));
            let obj = try!(new_repo.revparse_single(name));
            if let Ok(tag) = obj.into_tag() {
                let tag_name = try!(tag.name().ok_or(Error::from_str("Could not parse tag name"))).to_owned();
                new_tags.push((tag.target_id(), tag_name));
            }
        }
        Ok(GitJournal {
            repo: new_repo,
            tags: new_tags,
        })
    }

    fn parse_log(&self, revision_range: &str) -> Result<(), Error> {
        let mut revwalk = try!(self.repo.revwalk());
        revwalk.set_sorting(git2::SORT_TIME);

        // Fill the revwalk with the selected revisions.
        let revspec = try!(self.repo.revparse(&revision_range));
        if revspec.mode().contains(git2::REVPARSE_SINGLE) {
            try!(revwalk.push(revspec.from().unwrap().id()));
        } else {
            let from = revspec.from().unwrap().id();
            let to = revspec.to().unwrap().id();
            try!(revwalk.push(to));
            if revspec.mode().contains(git2::REVPARSE_MERGE_BASE) {
                let base = try!(self.repo.merge_base(from, to));
                let o = try!(self.repo.find_object(base, Some(ObjectType::Commit)));
                try!(revwalk.push(o.id()));
            }
            try!(revwalk.hide(from));
        }

        // Iterate over the git objects and process them.
        for id in revwalk {
            let oid = try!(id);
            let commit = try!(self.repo.find_commit(oid));
            for &(tag_id, ref tag_name) in self.tags.iter().filter(|tag| tag.0.as_bytes() == oid.as_bytes()) {
                println!("Found tag {} for commit {}.", tag_name, tag_id);
            }
            self.print_commit(commit);
        }
        Ok(())
    }

    fn print_commit(&self, commit: Commit) {
        println!("Commit: {}", commit.id());

        if commit.parents().len() > 1 {
            print!("Merge:");
            for id in commit.parent_ids() {
                print!(" {:.8}", id);
            }
            println!("");
        }

        let author = commit.author();
        println!("Author: {}\n", author);

        for line in String::from_utf8_lossy(commit.message_bytes()).lines() {
            println!("\t{}", line);
        }
        println!("\n");
    }
}

fn main() {
    let yaml = load_yaml!("cli.yaml");
    let matches = App::from_yaml(yaml).get_matches();

    let path = matches.value_of("path").expect("Could not parse 'path' parameter");
    let revision_range = matches.value_of("revision_range").expect("Could not parse 'revision range' parameter");

    match GitJournal::new(path) {
        Ok(journal) => journal.parse_log(revision_range).expect("Log parsing error"),
        Err(e) => println!("{}", e),
    }
}
<|endoftext|>


<commit_before>/*
 * Tool statically unpacking HotS binaries.
 *
 * The MIT License (MIT)
 *
 * Copyright (c) 2015 athre0z
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

extern crate crypto;
mod pe;

use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::error::Error;
use std::mem;
use std::slice;
use std::env;

use crypto::{buffer, aes, blockmodes};

#[allow(unused_imports)]
use crypto::symmetriccipher::Encryptor;

/// Entry-point.
fn main() {
    println!("HotS-unpack by athre0z (c) 2015 zyantific.com\n");

    // Parse arguments.
    let args: Vec<_> = env::args().collect();
    let (in_path, out_path) = match args.len() {
        2 => (&args[1], "out.exe".to_string()),
        3 => (&args[1], args[2].clone()),
        _ => {
            println!("[*] Usage: {} <input binary> [output binary]", args[0]);
            return;
        }
    };

    // Read file.
    let mut f = match File::open(&Path::new(&in_path)) {
        Err(why) => panic!("[-] Couldn't open input file: {}", Error::description(&why)),
        Ok(file) => file,
    };

    let mut data = Vec::new();
    let data_size = match f.read_to_end(&mut data) {
        Err(why) => panic!("[-] Couldn't read input file: {}", Error::description(&why)),
        Ok(x) => x
    };

    println!("[+] Read file: {} bytes", data_size);

    // Parse DOS-header.
    let mz = get_data::<pe::ImageDosHeader>(&mut data, 0);
    if unsafe { (*mz).magic } != pe::DOS_MAGIC {
        panic!("[-] Invalid input file: invalid DOS header");
    }

    // Parse NT-header.
    let nt_header_offs = unsafe { (*mz).lfanew } as usize;
    let nt = get_data::<pe::ImageNtHeaders32>(&mut data, nt_header_offs);
    if unsafe { (*nt).signature } != pe::PE_MAGIC {
        panic!("[-] Invalid input file: invalid NT header");
    }

    // Obtain useful information.
    let sec_table_offs = nt_header_offs + mem::size_of::<pe::ImageNtHeaders32>();
    let num_secs = unsafe { (*nt).file_header.number_of_sections } as usize;
    let data_offs = sec_table_offs
        + mem::size_of::<pe::ImageSectionHeader>() * num_secs;
    let image_base = unsafe { (*nt).opt_header.image_base } as usize;

    // Scan for chunk table.
    println!("[*] Heuristically locating chunk table ...");
    let mut chunk_table_offs = 0;
    let mut chunk_table_len = 0;

    'scan: for i in data_offs..data.len() {
        let mut offs = 0usize;
        let mut last_addr = 0u32;

        // Check if candidate looks like a chunk table.
        while i + offs + mem::size_of::<ChunkEntry>() <= data.len() {
            unsafe {
                let cur = get_data::<ChunkEntry>(&mut data, i + offs);

                // The chunk table's end is indicated with an 0xFFFFFFFF.
                if offs >= 1000 * mem::size_of::<ChunkEntry>() && (*cur).offset == 0xFFFFFFFF {
                    chunk_table_offs = i;
                    chunk_table_len = (offs / mem::size_of::<ChunkEntry>()) - 1;
                    break 'scan;
                }

                // The chunk table entries are stored ascending.
                if (*cur).offset <= last_addr {
                    break;
                }

                last_addr = (*cur).offset;
                offs += mem::size_of::<ChunkEntry>();
            }
        }
    }

    if chunk_table_offs == 0 || chunk_table_len == 0 {
        panic!("[-] Unable to locate chunk table.");
    }

    println!("[+] Found chunk table @ FO 0x{:x} ({} entries).",
        chunk_table_offs, chunk_table_len);

    // Scan for AES key.
    println!("[*] Locating AES key ...");

    // 0xFF -> wildcard
    let pattern = [
        0x68, 0x80, 0x00, 0x00, 0x00, // push 80h
        0x68, 0xFF, 0xFF, 0xFF, 0xFF, // push offset g_aesInputKey
        0x68, 0xFF, 0xFF, 0xFF, 0xFF, // push offset g_aesKey
        0xE8, 0xFF, 0xFF, 0xFF, 0xFF, // call AesSetDecryptKey
        0xE9,                         // jmp  XXX
    ];

    let mut match_ = None;
    for offs in sec_table_offs..(data.len() - pattern.len()) {
        let mut exhausted = true;
        for (i, cur_pat_byte) in pattern.iter().enumerate() {
            if *cur_pat_byte == 0xFF { continue; }
            if data[offs + i] != *cur_pat_byte {
                exhausted = false;
                break;
            }
        }

        if exhausted {
            match_ = Some(offs);
            break;
        }
    }

    let aes_key_rva = match match_ {
        Some(x) => unsafe { *get_data::<u32>(&mut data, x + 6) as usize - image_base },
        None => panic!("[-] Unable to locate AES key.")
    };

    print!("[+] Found AES key @ RVA 0x{:x}: ", aes_key_rva);

    // Dump AES key.
    let aes_key = unsafe { *get_data_rva::<[u8; 16]>(
        &mut data, sec_table_offs, num_secs, aes_key_rva) };

    for i in aes_key.iter() {
        print!("{:02x} ", i);
    }
    print!("\n");

    // Collect all encrypted chunks into one buffer (we got plenty of memory, don't we? :p)
    println!("[*] Merging encrypted chunks ...", );
    let mut encrypted_blob = Vec::<u8>::new();
    for i in 0..chunk_table_len {
        let chunk_info = get_data::<ChunkEntry>(
            &mut data, chunk_table_offs + mem::size_of::<ChunkEntry>() * i);
        unsafe {
            encrypted_blob.extend(get_slice_rva::<u8>(
                &mut data, sec_table_offs, num_secs, (*chunk_info).offset as usize,
                (*chunk_info).size as usize).to_vec());
        }
    }

    let padding_length = 16 - encrypted_blob.len() % 16;
    if padding_length > 0 {
        encrypted_blob.extend(vec![0; padding_length]);
    }

    // Decrypt blob.
    println!("[*] Decrypting ...");
    let mut decryptor = aes::ecb_decryptor(
        aes::KeySize::KeySize128,
        &aes_key,
        blockmodes::NoPadding);
    let mut decrypted_blob = vec![0u8; encrypted_blob.len()];

    {
        let mut read_buf = buffer::RefReadBuffer::new(&encrypted_blob);
        let mut write_buf = buffer::RefWriteBuffer::new(&mut decrypted_blob);
        match decryptor.decrypt(&mut read_buf, &mut write_buf, true) {
            Err(why) => panic!("[-] Unable to decrypt data: {:?}", why),
            Ok(result) => match result {
                buffer::BufferResult::BufferUnderflow => {},
                buffer::BufferResult::BufferOverflow => unreachable!(),
            }
        }
    }

    // Write data from blob back into our file buffer.
    println!("[*] Putting decrypted data where it belongs ...");
    let mut read_offset = 0usize;
    for i in 0..chunk_table_len {
        let chunk_info = get_data::<ChunkEntry>(
            &mut data, chunk_table_offs + mem::size_of::<ChunkEntry>() * i);
        unsafe {
            let chunk_size = (*chunk_info).size as usize;
            let mut dst = get_slice_rva::<u8>(&mut data, sec_table_offs, num_secs,
                (*chunk_info).offset as usize, chunk_size);
            let src = &decrypted_blob[read_offset..read_offset + chunk_size];

            for (cur_src_byte, cur_dst_byte) in src.iter().zip(dst.iter_mut()) {
                *cur_dst_byte = *cur_src_byte;
            }

            read_offset += chunk_size;
        }
    }

    // Blizz moved the IAT to the .reloc section, which isn't loaded by IDA
    // when loading a file without the "manual load" option, which results in
    // unresolved API calls. Rename it to make IDA map it into the IDB by default.
    println!("[*] Fixing section table ...");
    for i in 0..num_secs {
        let cur_sec = get_data::<pe::ImageSectionHeader>(
            &mut data, sec_table_offs + mem::size_of::<pe::ImageSectionHeader>() * i);

        unsafe {
            if (*cur_sec).name.iter().zip(".reloc".as_bytes().iter()).all(|(a, b)| a == b) {
                for (src, dst) in ".reloc_".as_bytes().iter().zip((*cur_sec).name.iter_mut()) {
                    *dst = *src;
                }
            }
        }
    }

    // Write output file.
    println!("[*] Writing output file ...");
    let mut out_f = match File::create(&Path::new(&out_path)) {
        Err(why) => panic!("[-] Couldn't open output file: {}", Error::description(&why)),
        Ok(file) => file,
    };

    match out_f.write_all(&data) {
        Err(why) => panic!("[-] Couldn't write to output file: {}", Error::description(&why)),
        Ok(..) => println!("[+] Wrote {:?} bytes ({}).", data.len(), &out_path),
    }
}

#[repr(C, packed)]
struct ChunkEntry {
    pub offset: u32,
    pub size: u32,
}

/// Obtains a ptr to a struct inside the input file by it's file offset.
fn get_data<T>(data: &mut Vec<u8>, file_offset: usize) -> *mut T {
    if mem::size_of::<T>() > data.len() - file_offset {
        panic!("[-] Invalid input file: file address exceeds file boundaries.");
    }
    unsafe { mem::transmute(data.as_ptr() as usize + file_offset) }
}

/// Obtains a ptr to a struct inside the input file by it's RVA.
fn get_data_rva<T>(data: &mut Vec<u8>, sec_table_offs: usize,
    num_secs: usize, rva: usize) -> *mut T {

    match rva_to_fo(data, sec_table_offs, num_secs, rva) {
        Some(rva) => get_data::<T>(data, rva),
        None => panic!("[-] Invalid input file: Pointer to an invalid RVA encountered.")
    }
}

/// Obtains a mutable slice from the input file by it's RVA.
fn get_slice_rva<'a, T>(data: &'a mut Vec<u8>, sec_table_offs: usize,
    num_secs: usize, rva: usize, len: usize) -> &'a mut [T] {

    let file_offset = match rva_to_fo(data, sec_table_offs, num_secs, rva) {
        Some(fo) => fo,
        None => panic!("[-] Invalid input file: Pointer to an invalid RVA encountered.")
    };

    if mem::size_of::<T>() * len > data.len() - file_offset {
        panic!("[-] Invalid input file: file address exceeds file boundaries.");
    }

    unsafe { slice::from_raw_parts_mut(mem::transmute(
        data.as_ptr() as usize + file_offset), len) }
}

/// Translates an RVA to a file offset.
fn rva_to_fo(data: &mut Vec<u8>, sec_table_offs: usize,
    num_secs: usize, rva: usize) -> Option<usize> {

    for i in 0..num_secs {
        let cur_sec = get_data::<pe::ImageSectionHeader>(
            data, sec_table_offs + mem::size_of::<pe::ImageSectionHeader>() * i);

        unsafe {
            let sec_va = (*cur_sec).virtual_address as usize;
            if rva >= sec_va && rva <= sec_va + (*cur_sec).virtual_size as usize {
                return Some(rva - sec_va + (*cur_sec).pointer_to_raw_data as usize);
            }
        }
    }

    None
}<commit_msg>refactored chunk table related code in order to ease implementation of new chunk table format<commit_after>/*
 * Tool statically unpacking HotS binaries.
 *
 * The MIT License (MIT)
 *
 * Copyright (c) 2015 athre0z
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

extern crate crypto;
mod pe;

use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::error::Error;
use std::mem;
use std::slice;
use std::env;

use crypto::{buffer, aes, blockmodes};

#[allow(unused_imports)]
use crypto::symmetriccipher::Encryptor;

/// Entry-point.
fn main() {
    println!("HotS-unpack by athre0z (c) 2015 zyantific.com\n");

    // Parse arguments.
    let args: Vec<_> = env::args().collect();
    let (in_path, out_path) = match args.len() {
        2 => (&args[1], "out.exe".to_string()),
        3 => (&args[1], args[2].clone()),
        _ => {
            println!("[*] Usage: {} <input binary> [output binary]", args[0]);
            return;
        }
    };

    // Read file.
    let mut f = match File::open(&Path::new(&in_path)) {
        Err(why) => panic!("[-] Couldn't open input file: {}", Error::description(&why)),
        Ok(file) => file,
    };

    let mut data = Vec::new();
    let data_size = match f.read_to_end(&mut data) {
        Err(why) => panic!("[-] Couldn't read input file: {}", Error::description(&why)),
        Ok(x) => x
    };

    println!("[+] Read file: {} bytes", data_size);

    // Parse DOS-header.
    let mz = get_data::<pe::ImageDosHeader>(&mut data, 0);
    if unsafe { (*mz).magic } != pe::DOS_MAGIC {
        panic!("[-] Invalid input file: invalid DOS header");
    }

    // Parse NT-header.
    let nt_header_offs = unsafe { (*mz).lfanew } as usize;
    let nt = get_data::<pe::ImageNtHeaders32>(&mut data, nt_header_offs);
    if unsafe { (*nt).signature } != pe::PE_MAGIC {
        panic!("[-] Invalid input file: invalid NT header");
    }

    // Obtain useful information.
    let sec_table_offs = nt_header_offs + mem::size_of::<pe::ImageNtHeaders32>();
    let num_secs = unsafe { (*nt).file_header.number_of_sections } as usize;
    let data_offs = sec_table_offs
        + mem::size_of::<pe::ImageSectionHeader>() * num_secs;
    let image_base = unsafe { (*nt).opt_header.image_base } as usize;

    let mut chunk_table_opt: Option<Vec<ChunkEntry>> = None;

    // Scan for chunk table (old format).
    {
        println!("[*] Heuristically locating chunk table ...");
        let mut chunk_table_offs = 0;
        let mut chunk_table_len = 0;

        'scan: for i in data_offs..data.len() {
            let mut offs = 0usize;
            let mut last_addr = 0u32;

            // Check if candidate looks like a chunk table.
            while i + offs + mem::size_of::<ChunkEntry>() <= data.len() {
                unsafe {
                    let cur = get_data::<ChunkEntry>(&mut data, i + offs);

                    // The chunk table's end is indicated with an 0xFFFFFFFF.
                    if offs >= 1000 * mem::size_of::<ChunkEntry>() && (*cur).offset == 0xFFFFFFFF {
                        chunk_table_offs = i;
                        chunk_table_len = (offs / mem::size_of::<ChunkEntry>()) - 1;
                        break 'scan;
                    }

                    // The chunk table entries are stored ascending.
                    if (*cur).offset <= last_addr {
                        break;
                    }

                    last_addr = (*cur).offset;
                    offs += mem::size_of::<ChunkEntry>();
                }
            }
        }

        if chunk_table_offs != 0 && chunk_table_len != 0 {
            unsafe {
                chunk_table_opt = Some(slice::from_raw_parts(get_data::<ChunkEntry>(
                    &mut data, chunk_table_offs), chunk_table_len).to_vec());
            }
        }
    }

    let chunk_table = match chunk_table_opt {
        None => panic!("[-] Unable to locate chunk table."),
        Some(x) => x,
    };

    println!("[+] Found chunk table.");

    // Scan for AES key.
    println!("[*] Locating AES key ...");

    // 0xFF -> wildcard
    let pattern = [
        0x68, 0x80, 0x00, 0x00, 0x00, // push 80h
        0x68, 0xFF, 0xFF, 0xFF, 0xFF, // push offset g_aesInputKey
        0x68, 0xFF, 0xFF, 0xFF, 0xFF, // push offset g_aesKey
        0xE8, 0xFF, 0xFF, 0xFF, 0xFF, // call AesSetDecryptKey
        0xE9,                         // jmp  XXX
    ];

    let mut match_ = None;
    for offs in sec_table_offs..(data.len() - pattern.len()) {
        let mut exhausted = true;
        for (i, cur_pat_byte) in pattern.iter().enumerate() {
            if *cur_pat_byte == 0xFF { continue; }
            if data[offs + i] != *cur_pat_byte {
                exhausted = false;
                break;
            }
        }

        if exhausted {
            match_ = Some(offs);
            break;
        }
    }

    let aes_key_rva = match match_ {
        Some(x) => unsafe { *get_data::<u32>(&mut data, x + 6) as usize - image_base },
        None => panic!("[-] Unable to locate AES key.")
    };

    print!("[+] Found AES key @ RVA 0x{:x}: ", aes_key_rva);

    // Dump AES key.
    let aes_key = unsafe { *get_data_rva::<[u8; 16]>(
        &mut data, sec_table_offs, num_secs, aes_key_rva) };

    for i in aes_key.iter() {
        print!("{:02x} ", i);
    }
    print!("\n");

    // Collect all encrypted chunks into one buffer (we got plenty of memory, don't we? :p)
    println!("[*] Merging encrypted chunks ...", );
    let mut encrypted_blob = Vec::<u8>::new();
    for cur_chunk in &chunk_table {
        encrypted_blob.extend(get_slice_rva::<u8>(
            &mut data, sec_table_offs, num_secs, cur_chunk.offset as usize,
            cur_chunk.size as usize).to_vec());
    }

    let padding_length = 16 - encrypted_blob.len() % 16;
    if padding_length > 0 {
        encrypted_blob.extend(vec![0; padding_length]);
    }

    // Decrypt blob.
    println!("[*] Decrypting ...");
    let mut decryptor = aes::ecb_decryptor(
        aes::KeySize::KeySize128,
        &aes_key,
        blockmodes::NoPadding);
    let mut decrypted_blob = vec![0u8; encrypted_blob.len()];

    {
        let mut read_buf = buffer::RefReadBuffer::new(&encrypted_blob);
        let mut write_buf = buffer::RefWriteBuffer::new(&mut decrypted_blob);
        match decryptor.decrypt(&mut read_buf, &mut write_buf, true) {
            Err(why) => panic!("[-] Unable to decrypt data: {:?}", why),
            Ok(result) => match result {
                buffer::BufferResult::BufferUnderflow => {},
                buffer::BufferResult::BufferOverflow => unreachable!(),
            }
        }
    }

    // Write data from blob back into our file buffer.
    println!("[*] Putting decrypted data where it belongs ...");
    let mut read_offset = 0usize;
    for cur_chunk in &chunk_table {
        let chunk_size = cur_chunk.size as usize;
        let mut dst = get_slice_rva::<u8>(&mut data, sec_table_offs, num_secs,
            cur_chunk.offset as usize, chunk_size);
        let src = &decrypted_blob[read_offset..read_offset + chunk_size];

        for (cur_src_byte, cur_dst_byte) in src.iter().zip(dst.iter_mut()) {
            *cur_dst_byte = *cur_src_byte;
        }

        read_offset += chunk_size;
    }

    // Blizz moved the IAT to the .reloc section, which isn't loaded by IDA
    // when loading a file without the "manual load" option, which results in
    // unresolved API calls. Rename it to make IDA map it into the IDB by default.
    println!("[*] Fixing section table ...");
    for i in 0..num_secs {
        let cur_sec = get_data::<pe::ImageSectionHeader>(
            &mut data, sec_table_offs + mem::size_of::<pe::ImageSectionHeader>() * i);

        unsafe {
            if (*cur_sec).name.iter().zip(".reloc".as_bytes().iter()).all(|(a, b)| a == b) {
                for (src, dst) in ".reloc_".as_bytes().iter().zip((*cur_sec).name.iter_mut()) {
                    *dst = *src;
                }
            }
        }
    }

    // Write output file.
    println!("[*] Writing output file ...");
    let mut out_f = match File::create(&Path::new(&out_path)) {
        Err(why) => panic!("[-] Couldn't open output file: {}", Error::description(&why)),
        Ok(file) => file,
    };

    match out_f.write_all(&data) {
        Err(why) => panic!("[-] Couldn't write to output file: {}", Error::description(&why)),
        Ok(..) => println!("[+] Wrote {:?} bytes ({}).", data.len(), &out_path),
    }
}

#[repr(C, packed)]
#[derive(Clone)]
struct ChunkEntry {
    pub offset: u32,
    pub size: u32,
}

/// Obtains a ptr to a struct inside the input file by it's file offset.
fn get_data<T>(data: &mut Vec<u8>, file_offset: usize) -> *mut T {
    if mem::size_of::<T>() > data.len() - file_offset {
        panic!("[-] Invalid input file: file address exceeds file boundaries.");
    }
    unsafe { mem::transmute(data.as_ptr() as usize + file_offset) }
}

/// Obtains a ptr to a struct inside the input file by it's RVA.
fn get_data_rva<T>(data: &mut Vec<u8>, sec_table_offs: usize,
    num_secs: usize, rva: usize) -> *mut T {

    match rva_to_fo(data, sec_table_offs, num_secs, rva) {
        Some(rva) => get_data::<T>(data, rva),
        None => panic!("[-] Invalid input file: Pointer to an invalid RVA encountered.")
    }
}

/// Obtains a mutable slice from the input file by it's RVA.
fn get_slice_rva<'a, T>(data: &'a mut Vec<u8>, sec_table_offs: usize,
    num_secs: usize, rva: usize, len: usize) -> &'a mut [T] {

    let file_offset = match rva_to_fo(data, sec_table_offs, num_secs, rva) {
        Some(fo) => fo,
        None => panic!("[-] Invalid input file: Pointer to an invalid RVA encountered.")
    };

    if mem::size_of::<T>() * len > data.len() - file_offset {
        panic!("[-] Invalid input file: file address exceeds file boundaries.");
    }

    unsafe { slice::from_raw_parts_mut(mem::transmute(
        data.as_ptr() as usize + file_offset), len) }
}

/// Translates an RVA to a file offset.
fn rva_to_fo(data: &mut Vec<u8>, sec_table_offs: usize,
    num_secs: usize, rva: usize) -> Option<usize> {

    for i in 0..num_secs {
        let cur_sec = get_data::<pe::ImageSectionHeader>(
            data, sec_table_offs + mem::size_of::<pe::ImageSectionHeader>() * i);

        unsafe {
            let sec_va = (*cur_sec).virtual_address as usize;
            if rva >= sec_va && rva <= sec_va + (*cur_sec).virtual_size as usize {
                return Some(rva - sec_va + (*cur_sec).pointer_to_raw_data as usize);
            }
        }
    }

    None
}<|endoftext|>


<commit_before>#![feature(duration_span)]

extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::time::Duration;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample, Colorf};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                } else {
                    block_samples.push(ImageSample::new(s.0, s.1, Colorf::black()));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, n: u32, spp: usize){
    let mut pool = scoped_threadpool::Pool::new(n);
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let r = &*rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let d = Duration::span(|| render_parallel(&mut rt, &scene, n, spp));
        let time = d.as_secs() as f64 + (d.subsec_nanos() as f64) / 1_000_000_000.0;
        println!("Frame {}: Rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: Rendered to '{}'\n--------------------", i, out_file.display());
    }
}

<commit_msg>USe a ref binding instead of the &* wierdness<commit_after>#![feature(duration_span)]

extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::time::Duration;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample, Colorf};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                } else {
                    block_samples.push(ImageSample::new(s.0, s.1, Colorf::black()));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, n: u32, spp: usize){
    let mut pool = scoped_threadpool::Pool::new(n);
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let ref r = rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let d = Duration::span(|| render_parallel(&mut rt, &scene, n, spp));
        let time = d.as_secs() as f64 + (d.subsec_nanos() as f64) / 1_000_000_000.0;
        println!("Frame {}: Rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: Rendered to '{}'\n--------------------", i, out_file.display());
    }
}

<|endoftext|>


<commit_before>extern crate tempfile;
extern crate rustyline;

mod lexer;
mod parser;
mod evaluator;
mod error;
mod constants;
mod function;

use rustyline::error::ReadlineError;
use rustyline::Editor;
use tempfile::NamedTempFile;
use constants::VERSION;
use lexer::lexer;
use parser::Parser;
use evaluator::Evaluator;
use error::{LexerError, ParserError, EvaluatorError};

fn prelude() {
    println!("bspl {}", VERSION);
    println!("Bit-Shift-Print Loop");
    println!("Type 'help', 'license', or 'version' for more information.");
}

fn error_message(width: usize, msg: &str) {
    println!("{caret:>width$}\n.. {}", msg, caret = "^", width = width);
}

fn display_results(results: Vec<String>) {
    if let Some((final_result, results)) = results.split_last(){
        for result in results {
            println!(".. {}", result);
        }

        let parsed_final_result = final_result.parse::<u32>().unwrap();
        println!("D: {}", parsed_final_result);
        println!("H: {:#x}", parsed_final_result);
        println!("B: {:#b}", parsed_final_result);
    }
}

fn repl() {
    let mut repl = Editor::<()>::new();
    let prompt = "=> ";
    let prompt_len = prompt.len() + 1;
    let tmp_file = NamedTempFile::new().unwrap();
    let _ = repl.load_history(tmp_file.path());

    let parser = Parser::default();
    let evaluator = Evaluator::default();

    loop {
        match repl.readline(prompt) {
            Ok(line) => {
                repl.add_history_entry(&line);
                match lexer(&line) {
                    Ok(tokens) => {
                        match parser.parse(tokens) {
                            Ok(parsed_tokens) => {
                                match evaluator.evaluate(parsed_tokens) {
                                    Ok(result) => {
                                        display_results(result);
                                    }
                                    Err(EvaluatorError::MissingArgument(position)) => {
                                        error_message(position + prompt_len,
                                                      "Missing argument from expression");
                                    }
                                    Err(EvaluatorError::TooManyArguments) => {
                                        error_message(prompt_len,
                                                      "Too many arguments in expression");
                                    }
                                    Err(EvaluatorError::OverflowShift(position)) => {
                                        error_message(position + prompt_len,
                                                      "Shift overflow error");
                                    }
                                    Err(EvaluatorError::UnknownKeyword(position)) => {
                                        error_message(position + prompt_len, "Not a valid keyword");
                                    }
                                    Err(EvaluatorError::Exit) => break,
                                }
                            }
                            Err(ParserError::MissingOpeningBracket(position)) => {
                                error_message(position + prompt_len, "Missing an opening bracket");
                            }
                            Err(ParserError::MissingClosingBracket(position)) => {
                                error_message(position + prompt_len, "Missing a closing bracket");
                            }
                            Err(ParserError::KeywordError(position)) => {
                                error_message(position + prompt_len,
                                              "Cannot use keyword in expression");
                            }
                        }
                    }
                    Err(LexerError::RadixError(position)) => {
                        error_message(position + prompt_len,
                                      "Not a valid decimal, hexadecimal, or keyword");
                    }
                    Err(LexerError::UnknownOperator(position)) => {
                        error_message(position + prompt_len, "Not a valid operator");
                    }
                }
            }
            Err(ReadlineError::Eof) => break,
            Err(ReadlineError::Interrupted) => {
                println!("Type 'exit' or press Ctrl-D to leave bspl");
                continue;
            }
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }

        }
        repl.save_history(tmp_file.path()).unwrap();
    }
}

pub fn main() {
    prelude();
    repl();
}
<commit_msg>Check final result can be parsed as u32<commit_after>extern crate tempfile;
extern crate rustyline;

mod lexer;
mod parser;
mod evaluator;
mod error;
mod constants;
mod function;

use rustyline::error::ReadlineError;
use rustyline::Editor;
use tempfile::NamedTempFile;
use constants::VERSION;
use lexer::lexer;
use parser::Parser;
use evaluator::Evaluator;
use error::{LexerError, ParserError, EvaluatorError};

fn prelude() {
    println!("bspl {}", VERSION);
    println!("Bit-Shift-Print Loop");
    println!("Type 'help', 'license', or 'version' for more information.");
}

fn error_message(width: usize, msg: &str) {
    println!("{caret:>width$}\n.. {}", msg, caret = "^", width = width);
}

fn display_results(results: Vec<String>) {
    if let Some((final_result, results)) = results.split_last(){
        for result in results {
            println!(".. {}", result);
        }

        if let Ok(parsed_final_result) = final_result.parse::<u32>(){
            println!("D: {}", parsed_final_result);
            println!("H: {:#x}", parsed_final_result);
            println!("B: {:#b}", parsed_final_result);
        } else {
            println!(".. {}", final_result);
        }
    }
}

fn repl() {
    let mut repl = Editor::<()>::new();
    let prompt = "=> ";
    let prompt_len = prompt.len() + 1;
    let tmp_file = NamedTempFile::new().unwrap();
    let _ = repl.load_history(tmp_file.path());

    let parser = Parser::default();
    let evaluator = Evaluator::default();

    loop {
        match repl.readline(prompt) {
            Ok(line) => {
                repl.add_history_entry(&line);
                match lexer(&line) {
                    Ok(tokens) => {
                        match parser.parse(tokens) {
                            Ok(parsed_tokens) => {
                                match evaluator.evaluate(parsed_tokens) {
                                    Ok(result) => {
                                        display_results(result);
                                    }
                                    Err(EvaluatorError::MissingArgument(position)) => {
                                        error_message(position + prompt_len,
                                                      "Missing argument from expression");
                                    }
                                    Err(EvaluatorError::TooManyArguments) => {
                                        error_message(prompt_len,
                                                      "Too many arguments in expression");
                                    }
                                    Err(EvaluatorError::OverflowShift(position)) => {
                                        error_message(position + prompt_len,
                                                      "Shift overflow error");
                                    }
                                    Err(EvaluatorError::UnknownKeyword(position)) => {
                                        error_message(position + prompt_len, "Not a valid keyword");
                                    }
                                    Err(EvaluatorError::Exit) => break,
                                }
                            }
                            Err(ParserError::MissingOpeningBracket(position)) => {
                                error_message(position + prompt_len, "Missing an opening bracket");
                            }
                            Err(ParserError::MissingClosingBracket(position)) => {
                                error_message(position + prompt_len, "Missing a closing bracket");
                            }
                            Err(ParserError::KeywordError(position)) => {
                                error_message(position + prompt_len,
                                              "Cannot use keyword in expression");
                            }
                        }
                    }
                    Err(LexerError::RadixError(position)) => {
                        error_message(position + prompt_len,
                                      "Not a valid decimal, hexadecimal, or keyword");
                    }
                    Err(LexerError::UnknownOperator(position)) => {
                        error_message(position + prompt_len, "Not a valid operator");
                    }
                }
            }
            Err(ReadlineError::Eof) => break,
            Err(ReadlineError::Interrupted) => {
                println!("Type 'exit' or press Ctrl-D to leave bspl");
                continue;
            }
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }

        }
        repl.save_history(tmp_file.path()).unwrap();
    }
}

pub fn main() {
    prelude();
    repl();
}
<|endoftext|>