<commit_before>/*
 * Copyright (c) 2016-2018 Boucher, Antoni <bouanto@zoho.com>
 *
 * 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.
 */

use std::collections::HashMap;

use glib::Cast;
use gtk::{
    self,
    Container,
    ContainerExt,
    FlowBoxExt,
    IsA,
    SelectionMode,
    WidgetExt,
};
use relm::{
    Component,
    ContainerWidget,
    Relm,
    Widget,
    timeout,
};
use relm_attributes::widget;
use webkit2gtk::{
    Download,
    DownloadExt,
    Error,
};
use webkit2gtk::DownloadError::CancelledByUser;

use download_view::DownloadView;
use download_view::Msg::{
    Cancel,
    Destination,
    DownloadError,
    OriginalDestination,
    Remove,
    SetToOpen,
};
use self::Msg::*;

const DOWNLOAD_TIME_BEFORE_HIDE: u32 = 2;

pub struct Model {
    download_count: u32,
    download_views: HashMap<Download, Component<DownloadView>>,
    relm: Relm<DownloadListView>,
}

#[derive(Msg)]
pub enum Msg {
    ActiveDownloads(bool),
    Add(Download),
    AddFileToOpen(Download),
    DelayedRemove(Download),
    DownloadCancel(Download),
    DownloadDestination(Download, String),
    DownloadFailed(Error, Download),
    DownloadFinished(Download),
    DownloadListError(String),
    DownloadOriginalDestination(Download, String),
    DownloadRemove(Download),
}

#[widget]
impl Widget for DownloadListView {
    fn model(relm: &Relm<Self>, (): ()) -> Model {
        Model {
            download_count: 0,
            download_views: HashMap::new(),
            relm: relm.clone(),
        }
    }

    fn update(&mut self, event: Msg) {
        match event {
            // To be listened by the user.
            ActiveDownloads(_) => (),
            Add(download) => self.add(download),
            AddFileToOpen(download) => self.add_file_to_open(download),
            DelayedRemove(download) => self.delayed_remove(download),
            DownloadCancel(download) => self.download_cancel(download),
            DownloadDestination(download, destination) => self.download_destination(download, destination),
            DownloadFailed(error, download) => self.handle_failed(&error, download),
            DownloadFinished(_) => self.handle_finished(),
            DownloadListError(_) => (), // To be listened by the user.
            DownloadOriginalDestination(download, destination) =>
                self.download_original_destination(download, destination),
            DownloadRemove(download) => self.delete(download),
        }
    }

    view! {
        #[name="view"]
        gtk::FlowBox {
            selection_mode: SelectionMode::None,
        }
    }
}

impl DownloadListView {
    /// Add a new download.
    pub fn add(&mut self, download: Download) {
        self.model.download_count += 1;
        self.model.relm.stream().emit(ActiveDownloads(true));

        connect!(self.model.relm, download, connect_failed(download, error),
            DownloadFailed(error.clone(), download.clone()));
        connect!(self.model.relm, download, connect_finished(download), DownloadFinished(download.clone()));

        let download_view = self.view.add_widget::<DownloadView>(download.clone());
        if let Some(flow_child) = self.view.get_children().last() {
            flow_child.set_can_focus(false);
        }
        let down = download.clone();
        connect!(download_view@DownloadError(ref error), self.model.relm, DownloadListError(error.clone()));
        connect!(download_view@Remove, self.model.relm, DelayedRemove(down.clone()));

        // It is necessary to keep the download views because they are connected to events.
        self.model.download_views.insert(download, download_view);
    }

    /// Add a file to be opened when its download finish.
    pub fn add_file_to_open(&mut self, download: Download) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(SetToOpen);
        }
    }

    fn delayed_remove(&self, download: Download) {
        // Delete the view after a certain amount of time after the download finishes.
        timeout(self.model.relm.stream(), DOWNLOAD_TIME_BEFORE_HIDE, move || DownloadRemove(download.clone()));
    }

    /// Delete a view and remove it from its parent.
    fn delete(&mut self, download: Download) {
        if let Some(download_view) = self.model.download_views.remove(&download) {
            remove_from_flow_box(download_view.widget());
        }
    }

    fn download_cancel(&self, download: Download) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(Cancel);
        }
        // TODO: warning?
    }

    fn download_destination(&self, download: Download, destination: String) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(Destination(destination));
        }
        // TODO: warning?
    }

    fn download_original_destination(&self, download: Download, destination: String) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(OriginalDestination(destination));
        }
        // TODO: warning?
    }

    /// Handle the download failed event.
    /// Delete the view if the download was cancelled by the user.
    fn handle_failed(&mut self, error: &Error, download: Download) {
        warn!("Download failed: {}", error);
        if let Some(error) = error.kind::<::webkit2gtk::DownloadError>() {
            if error == CancelledByUser {
                self.delete(download);
            }
        }
    }

    /// Handle the download fisished event.
    fn handle_finished(&mut self) {
        self.model.download_count -= 1;

        if self.model.download_count == 0 {
            self.model.relm.stream().emit(ActiveDownloads(false));
        }
    }
}

/// Remove the progress bar from its `FlowBox` parent.
fn remove_from_flow_box<W: IsA<gtk::Widget> + WidgetExt>(widget: &W) {
    let child: Option<Container> = widget.get_parent()
        .and_then(|parent| parent.downcast().ok());
    // FlowBox children are wrapped inside FlowBoxChild, so we need to destroy this
    // FlowBoxChild (which is the parent of the widget) in order to remove it from
    // the FlowBox.
    if let Some(child) = child {
        child.destroy();
    }
}
<commit_msg>Fix timeout duration<commit_after>/*
 * Copyright (c) 2016-2018 Boucher, Antoni <bouanto@zoho.com>
 *
 * 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.
 */

use std::collections::HashMap;

use glib::Cast;
use gtk::{
    self,
    Container,
    ContainerExt,
    FlowBoxExt,
    IsA,
    SelectionMode,
    WidgetExt,
};
use relm::{
    Component,
    ContainerWidget,
    Relm,
    Widget,
    timeout,
};
use relm_attributes::widget;
use webkit2gtk::{
    Download,
    DownloadExt,
    Error,
};
use webkit2gtk::DownloadError::CancelledByUser;

use download_view::DownloadView;
use download_view::Msg::{
    Cancel,
    Destination,
    DownloadError,
    OriginalDestination,
    Remove,
    SetToOpen,
};
use self::Msg::*;

const DOWNLOAD_TIME_BEFORE_HIDE: u32 = 2000;

pub struct Model {
    download_count: u32,
    download_views: HashMap<Download, Component<DownloadView>>,
    relm: Relm<DownloadListView>,
}

#[derive(Msg)]
pub enum Msg {
    ActiveDownloads(bool),
    Add(Download),
    AddFileToOpen(Download),
    DelayedRemove(Download),
    DownloadCancel(Download),
    DownloadDestination(Download, String),
    DownloadFailed(Error, Download),
    DownloadFinished(Download),
    DownloadListError(String),
    DownloadOriginalDestination(Download, String),
    DownloadRemove(Download),
}

#[widget]
impl Widget for DownloadListView {
    fn model(relm: &Relm<Self>, (): ()) -> Model {
        Model {
            download_count: 0,
            download_views: HashMap::new(),
            relm: relm.clone(),
        }
    }

    fn update(&mut self, event: Msg) {
        match event {
            // To be listened by the user.
            ActiveDownloads(_) => (),
            Add(download) => self.add(download),
            AddFileToOpen(download) => self.add_file_to_open(download),
            DelayedRemove(download) => self.delayed_remove(download),
            DownloadCancel(download) => self.download_cancel(download),
            DownloadDestination(download, destination) => self.download_destination(download, destination),
            DownloadFailed(error, download) => self.handle_failed(&error, download),
            DownloadFinished(_) => self.handle_finished(),
            DownloadListError(_) => (), // To be listened by the user.
            DownloadOriginalDestination(download, destination) =>
                self.download_original_destination(download, destination),
            DownloadRemove(download) => self.delete(download),
        }
    }

    view! {
        #[name="view"]
        gtk::FlowBox {
            selection_mode: SelectionMode::None,
        }
    }
}

impl DownloadListView {
    /// Add a new download.
    pub fn add(&mut self, download: Download) {
        self.model.download_count += 1;
        self.model.relm.stream().emit(ActiveDownloads(true));

        connect!(self.model.relm, download, connect_failed(download, error),
            DownloadFailed(error.clone(), download.clone()));
        connect!(self.model.relm, download, connect_finished(download), DownloadFinished(download.clone()));

        let download_view = self.view.add_widget::<DownloadView>(download.clone());
        if let Some(flow_child) = self.view.get_children().last() {
            flow_child.set_can_focus(false);
        }
        let down = download.clone();
        connect!(download_view@DownloadError(ref error), self.model.relm, DownloadListError(error.clone()));
        connect!(download_view@Remove, self.model.relm, DelayedRemove(down.clone()));

        // It is necessary to keep the download views because they are connected to events.
        self.model.download_views.insert(download, download_view);
    }

    /// Add a file to be opened when its download finish.
    pub fn add_file_to_open(&mut self, download: Download) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(SetToOpen);
        }
    }

    fn delayed_remove(&self, download: Download) {
        // Delete the view after a certain amount of time after the download finishes.
        timeout(self.model.relm.stream(), DOWNLOAD_TIME_BEFORE_HIDE, move || DownloadRemove(download.clone()));
    }

    /// Delete a view and remove it from its parent.
    fn delete(&mut self, download: Download) {
        if let Some(download_view) = self.model.download_views.remove(&download) {
            remove_from_flow_box(download_view.widget());
        }
    }

    fn download_cancel(&self, download: Download) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(Cancel);
        }
        // TODO: warning?
    }

    fn download_destination(&self, download: Download, destination: String) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(Destination(destination));
        }
        // TODO: warning?
    }

    fn download_original_destination(&self, download: Download, destination: String) {
        if let Some(download_view) = self.model.download_views.get(&download) {
            download_view.emit(OriginalDestination(destination));
        }
        // TODO: warning?
    }

    /// Handle the download failed event.
    /// Delete the view if the download was cancelled by the user.
    fn handle_failed(&mut self, error: &Error, download: Download) {
        warn!("Download failed: {}", error);
        if let Some(error) = error.kind::<::webkit2gtk::DownloadError>() {
            if error == CancelledByUser {
                self.delete(download);
            }
        }
    }

    /// Handle the download fisished event.
    fn handle_finished(&mut self) {
        self.model.download_count -= 1;

        if self.model.download_count == 0 {
            self.model.relm.stream().emit(ActiveDownloads(false));
        }
    }
}

/// Remove the progress bar from its `FlowBox` parent.
fn remove_from_flow_box<W: IsA<gtk::Widget> + WidgetExt>(widget: &W) {
    let child: Option<Container> = widget.get_parent()
        .and_then(|parent| parent.downcast().ok());
    // FlowBox children are wrapped inside FlowBoxChild, so we need to destroy this
    // FlowBoxChild (which is the parent of the widget) in order to remove it from
    // the FlowBox.
    if let Some(child) = child {
        child.destroy();
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets using the /dev/bpf device.

use crate::bindings::bpf;
use crate::{DataLinkReceiver, DataLinkSender, NetworkInterface};

use pnet_sys;

use std::collections::VecDeque;
use std::ffi::CString;
use std::io;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::time::Duration;

static ETHERNET_HEADER_SIZE: usize = 14;

/// The BPF-specific configuration.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096.
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096.
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// The number of /dev/bpf* file descriptors to attempt before failing.
    ///
    /// This setting is only used on OS X - FreeBSD uses a single /dev/bpf rather than creating a
    /// new descriptor each time one is opened.
    ///
    /// Defaults to: 1000.
    pub bpf_fd_attempts: usize,
}

impl<'a> From<&'a super::Config> for Config {
    fn from(config: &super::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
            bpf_fd_attempts: config.bpf_fd_attempts,
            read_timeout: config.read_timeout,
            write_timeout: config.write_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            bpf_fd_attempts: 1000,
            read_timeout: None,
            write_timeout: None,
        }
    }
}

/// Create a datalink channel using the /dev/bpf device
// NOTE buffer must be word aligned.
#[inline]
pub fn channel(network_interface: &NetworkInterface, config: Config) -> io::Result<super::Channel> {
    #[cfg(any(target_os = "freebsd", target_os = "netbsd"))]
    fn get_fd(_attempts: usize) -> libc::c_int {
        unsafe {
            libc::open(
                CString::new(&b"/dev/bpf"[..]).unwrap().as_ptr(),
                libc::O_RDWR,
                0,
            )
        }
    }

    #[cfg(any(target_os = "openbsd", target_os = "macos", target_os = "ios"))]
    fn get_fd(attempts: usize) -> libc::c_int {
        for i in 0..attempts {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                let c_file_name = CString::new(file_name.as_bytes()).unwrap();
                libc::open(
                    c_file_name.as_ptr(),
                    libc::O_RDWR,
                    0,
                )
            };
            if fd != -1 {
                return fd;
            }
        }

        -1
    }

    #[cfg(any(target_os = "freebsd", target_os = "netbsd"))]
    fn set_feedback(fd: libc::c_int) -> io::Result<()> {
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(any(target_os = "macos", target_os = "openbsd", target_os = "ios"))]
    fn set_feedback(_fd: libc::c_int) -> io::Result<()> {
        Ok(())
    }

    let fd = get_fd(config.bpf_fd_attempts);
    if fd == -1 {
        return Err(io::Error::last_os_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    for (i, c) in network_interface.name.bytes().enumerate() {
        iface.ifr_name[i] = c as i8;
    }

    let buflen = config.read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the
    // buffer
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &1) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    let mut loopback = false;
    let mut allocated_read_buffer_size = config.read_buffer_size;
    // The loopback device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        loopback = true;
        // So we can guaranatee that we can have a header before the packet.
        // Loopback packets arrive without the header.
        allocated_read_buffer_size += ETHERNET_HEADER_SIZE;

        // Allow packets to be read back after they are written
        if let Err(e) = set_feedback(fd) {
            return Err(e);
        }
    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
    }

    // Enable nonblocking
    if unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            pnet_sys::close(fd);
        }
        return Err(err);
    }

    let fd = Arc::new(pnet_sys::FileDesc { fd: fd });
    let mut sender = Box::new(DataLinkSenderImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        write_buffer: vec![0; config.write_buffer_size],
        loopback: loopback,
        timeout: config
            .write_timeout
            .map(|to| pnet_sys::duration_to_timespec(to)),
    });
    unsafe {
        libc::FD_ZERO(&mut sender.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut sender.fd_set as *mut libc::fd_set);
    }
    let mut receiver = Box::new(DataLinkReceiverImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        read_buffer: vec![0; allocated_read_buffer_size],
        loopback: loopback,
        timeout: config
            .read_timeout
            .map(|to| pnet_sys::duration_to_timespec(to)),
        // Enough room for minimally sized packets without reallocating
        packets: VecDeque::with_capacity(allocated_read_buffer_size / 64),
    });
    unsafe {
        libc::FD_ZERO(&mut receiver.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut receiver.fd_set as *mut libc::fd_set);
    }

    Ok(super::Channel::Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    write_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(
        &mut self,
        num_packets: usize,
        packet_size: usize,
        func: &mut dyn FnMut(&mut [u8]),
    ) -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= self.write_buffer.len() {
            None
        } else {
            // If we're sending on the loopback device, discard the ethernet header.
            // The OS will prepend the packet with 4 bytes set to AF_INET.
            let offset = if self.loopback {
                ETHERNET_HEADER_SIZE
            } else {
                0
            };
            for chunk in self.write_buffer[..len].chunks_mut(packet_size) {
                func(chunk);
                let ret = unsafe {
                    libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                    libc::pselect(
                        self.fd.fd + 1,
                        ptr::null_mut(),
                        &mut self.fd_set as *mut libc::fd_set,
                        ptr::null_mut(),
                        self.timeout
                            .as_ref()
                            .map(|to| to as *const libc::timespec)
                            .unwrap_or(ptr::null()),
                        ptr::null(),
                    )
                };
                if ret == -1 {
                    // Error occured!
                    return Some(Err(io::Error::last_os_error()));
                } else if ret == 0 {
                    return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
                } else {
                    match unsafe {
                        libc::write(
                            self.fd.fd,
                            chunk.as_ptr().offset(offset as isize) as *const libc::c_void,
                            (chunk.len() - offset) as libc::size_t,
                        )
                    } {
                        len if len == -1 => return Some(Err(io::Error::last_os_error())),
                        _ => (),
                    }
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self, packet: &[u8], _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        // If we're sending on the loopback device, discard the ethernet header.
        // The OS will prepend the packet with 4 bytes set to AF_INET.
        let offset = if self.loopback {
            ETHERNET_HEADER_SIZE
        } else {
            0
        };
        let ret = unsafe {
            libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
            libc::pselect(
                self.fd.fd + 1,
                ptr::null_mut(),
                &mut self.fd_set as *mut libc::fd_set,
                ptr::null_mut(),
                self.timeout
                    .as_ref()
                    .map(|to| to as *const libc::timespec)
                    .unwrap_or(ptr::null()),
                ptr::null(),
            )
        };
        if ret == -1 {
            // Error occured!
            return Some(Err(io::Error::last_os_error()));
        } else if ret == 0 {
            return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
        } else {
            match unsafe {
                libc::write(
                    self.fd.fd,
                    packet.as_ptr().offset(offset as isize) as *const libc::c_void,
                    (packet.len() - offset) as libc::size_t,
                )
            } {
                len if len == -1 => Some(Err(io::Error::last_os_error())),
                _ => Some(Ok(())),
            }
        }
    }
}

struct DataLinkReceiverImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    read_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
    packets: VecDeque<(usize, usize)>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn next(&mut self) -> io::Result<&[u8]> {
        // Loopback packets arrive with a 4 byte header instead of normal ethernet header.
        // Discard that header and replace with zeroed out ethernet header.
        let (header_size, buffer_offset) = if self.loopback {
            (4, ETHERNET_HEADER_SIZE)
        } else {
            (0, 0)
        };
        if self.packets.is_empty() {
            let buffer = &mut self.read_buffer[buffer_offset..];
            let ret = unsafe {
                libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                libc::pselect(
                    self.fd.fd + 1,
                    &mut self.fd_set as *mut libc::fd_set,
                    ptr::null_mut(),
                    ptr::null_mut(),
                    self.timeout
                        .as_ref()
                        .map(|to| to as *const libc::timespec)
                        .unwrap_or(ptr::null()),
                    ptr::null(),
                )
            };
            if ret == -1 {
                return Err(io::Error::last_os_error());
            } else if ret == 0 {
                return Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"));
            } else {
                let buflen = match unsafe {
                    libc::read(
                        self.fd.fd,
                        buffer.as_ptr() as *mut libc::c_void,
                        buffer.len() as libc::size_t,
                    )
                } {
                    len if len > 0 => len,
                    _ => return Err(io::Error::last_os_error()),
                };
                let mut ptr = buffer.as_mut_ptr();
                let end = unsafe { buffer.as_ptr().offset(buflen as isize) };
                while (ptr as *const u8) < end {
                    unsafe {
                        let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                        let start =
                            ptr as isize + (*packet).bh_hdrlen as isize - buffer.as_ptr() as isize;
                        self.packets.push_back((
                            start as usize + header_size,
                            (*packet).bh_caplen as usize - header_size,
                        ));
                        let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                        ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                    }
                }
            }
        }
        let (start, mut len) = self.packets.pop_front().unwrap();
        len += buffer_offset;
        // Zero out part that will become fake ethernet header if on loopback.
        for i in (&mut self.read_buffer[start..start + buffer_offset]).iter_mut() {
            *i = 0;
        }
        Ok(&self.read_buffer[start..start + len])
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    #[path = "unix_interfaces.rs"]
    mod interfaces;
    interfaces::interfaces()
}
<commit_msg>Fix build on architectures with unsigned char<commit_after>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets using the /dev/bpf device.

use crate::bindings::bpf;
use crate::{DataLinkReceiver, DataLinkSender, NetworkInterface};

use pnet_sys;

use std::collections::VecDeque;
use std::ffi::CString;
use std::io;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::time::Duration;

static ETHERNET_HEADER_SIZE: usize = 14;

/// The BPF-specific configuration.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096.
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096.
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// The number of /dev/bpf* file descriptors to attempt before failing.
    ///
    /// This setting is only used on OS X - FreeBSD uses a single /dev/bpf rather than creating a
    /// new descriptor each time one is opened.
    ///
    /// Defaults to: 1000.
    pub bpf_fd_attempts: usize,
}

impl<'a> From<&'a super::Config> for Config {
    fn from(config: &super::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
            bpf_fd_attempts: config.bpf_fd_attempts,
            read_timeout: config.read_timeout,
            write_timeout: config.write_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            bpf_fd_attempts: 1000,
            read_timeout: None,
            write_timeout: None,
        }
    }
}

/// Create a datalink channel using the /dev/bpf device
// NOTE buffer must be word aligned.
#[inline]
pub fn channel(network_interface: &NetworkInterface, config: Config) -> io::Result<super::Channel> {
    #[cfg(any(target_os = "freebsd", target_os = "netbsd"))]
    fn get_fd(_attempts: usize) -> libc::c_int {
        unsafe {
            libc::open(
                CString::new(&b"/dev/bpf"[..]).unwrap().as_ptr(),
                libc::O_RDWR,
                0,
            )
        }
    }

    #[cfg(any(target_os = "openbsd", target_os = "macos", target_os = "ios"))]
    fn get_fd(attempts: usize) -> libc::c_int {
        for i in 0..attempts {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                let c_file_name = CString::new(file_name.as_bytes()).unwrap();
                libc::open(
                    c_file_name.as_ptr(),
                    libc::O_RDWR,
                    0,
                )
            };
            if fd != -1 {
                return fd;
            }
        }

        -1
    }

    #[cfg(any(target_os = "freebsd", target_os = "netbsd"))]
    fn set_feedback(fd: libc::c_int) -> io::Result<()> {
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(any(target_os = "macos", target_os = "openbsd", target_os = "ios"))]
    fn set_feedback(_fd: libc::c_int) -> io::Result<()> {
        Ok(())
    }

    let fd = get_fd(config.bpf_fd_attempts);
    if fd == -1 {
        return Err(io::Error::last_os_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    for (i, c) in network_interface.name.bytes().enumerate() {
        iface.ifr_name[i] = c as libc::c_char;
    }

    let buflen = config.read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the
    // buffer
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &1) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    let mut loopback = false;
    let mut allocated_read_buffer_size = config.read_buffer_size;
    // The loopback device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        loopback = true;
        // So we can guaranatee that we can have a header before the packet.
        // Loopback packets arrive without the header.
        allocated_read_buffer_size += ETHERNET_HEADER_SIZE;

        // Allow packets to be read back after they are written
        if let Err(e) = set_feedback(fd) {
            return Err(e);
        }
    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &1) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
    }

    // Enable nonblocking
    if unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            pnet_sys::close(fd);
        }
        return Err(err);
    }

    let fd = Arc::new(pnet_sys::FileDesc { fd: fd });
    let mut sender = Box::new(DataLinkSenderImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        write_buffer: vec![0; config.write_buffer_size],
        loopback: loopback,
        timeout: config
            .write_timeout
            .map(|to| pnet_sys::duration_to_timespec(to)),
    });
    unsafe {
        libc::FD_ZERO(&mut sender.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut sender.fd_set as *mut libc::fd_set);
    }
    let mut receiver = Box::new(DataLinkReceiverImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        read_buffer: vec![0; allocated_read_buffer_size],
        loopback: loopback,
        timeout: config
            .read_timeout
            .map(|to| pnet_sys::duration_to_timespec(to)),
        // Enough room for minimally sized packets without reallocating
        packets: VecDeque::with_capacity(allocated_read_buffer_size / 64),
    });
    unsafe {
        libc::FD_ZERO(&mut receiver.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut receiver.fd_set as *mut libc::fd_set);
    }

    Ok(super::Channel::Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    write_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(
        &mut self,
        num_packets: usize,
        packet_size: usize,
        func: &mut dyn FnMut(&mut [u8]),
    ) -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= self.write_buffer.len() {
            None
        } else {
            // If we're sending on the loopback device, discard the ethernet header.
            // The OS will prepend the packet with 4 bytes set to AF_INET.
            let offset = if self.loopback {
                ETHERNET_HEADER_SIZE
            } else {
                0
            };
            for chunk in self.write_buffer[..len].chunks_mut(packet_size) {
                func(chunk);
                let ret = unsafe {
                    libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                    libc::pselect(
                        self.fd.fd + 1,
                        ptr::null_mut(),
                        &mut self.fd_set as *mut libc::fd_set,
                        ptr::null_mut(),
                        self.timeout
                            .as_ref()
                            .map(|to| to as *const libc::timespec)
                            .unwrap_or(ptr::null()),
                        ptr::null(),
                    )
                };
                if ret == -1 {
                    // Error occured!
                    return Some(Err(io::Error::last_os_error()));
                } else if ret == 0 {
                    return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
                } else {
                    match unsafe {
                        libc::write(
                            self.fd.fd,
                            chunk.as_ptr().offset(offset as isize) as *const libc::c_void,
                            (chunk.len() - offset) as libc::size_t,
                        )
                    } {
                        len if len == -1 => return Some(Err(io::Error::last_os_error())),
                        _ => (),
                    }
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self, packet: &[u8], _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        // If we're sending on the loopback device, discard the ethernet header.
        // The OS will prepend the packet with 4 bytes set to AF_INET.
        let offset = if self.loopback {
            ETHERNET_HEADER_SIZE
        } else {
            0
        };
        let ret = unsafe {
            libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
            libc::pselect(
                self.fd.fd + 1,
                ptr::null_mut(),
                &mut self.fd_set as *mut libc::fd_set,
                ptr::null_mut(),
                self.timeout
                    .as_ref()
                    .map(|to| to as *const libc::timespec)
                    .unwrap_or(ptr::null()),
                ptr::null(),
            )
        };
        if ret == -1 {
            // Error occured!
            return Some(Err(io::Error::last_os_error()));
        } else if ret == 0 {
            return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
        } else {
            match unsafe {
                libc::write(
                    self.fd.fd,
                    packet.as_ptr().offset(offset as isize) as *const libc::c_void,
                    (packet.len() - offset) as libc::size_t,
                )
            } {
                len if len == -1 => Some(Err(io::Error::last_os_error())),
                _ => Some(Ok(())),
            }
        }
    }
}

struct DataLinkReceiverImpl {
    fd: Arc<pnet_sys::FileDesc>,
    fd_set: libc::fd_set,
    read_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
    packets: VecDeque<(usize, usize)>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn next(&mut self) -> io::Result<&[u8]> {
        // Loopback packets arrive with a 4 byte header instead of normal ethernet header.
        // Discard that header and replace with zeroed out ethernet header.
        let (header_size, buffer_offset) = if self.loopback {
            (4, ETHERNET_HEADER_SIZE)
        } else {
            (0, 0)
        };
        if self.packets.is_empty() {
            let buffer = &mut self.read_buffer[buffer_offset..];
            let ret = unsafe {
                libc::FD_SET(self.fd.fd, &mut self.fd_set as *mut libc::fd_set);
                libc::pselect(
                    self.fd.fd + 1,
                    &mut self.fd_set as *mut libc::fd_set,
                    ptr::null_mut(),
                    ptr::null_mut(),
                    self.timeout
                        .as_ref()
                        .map(|to| to as *const libc::timespec)
                        .unwrap_or(ptr::null()),
                    ptr::null(),
                )
            };
            if ret == -1 {
                return Err(io::Error::last_os_error());
            } else if ret == 0 {
                return Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"));
            } else {
                let buflen = match unsafe {
                    libc::read(
                        self.fd.fd,
                        buffer.as_ptr() as *mut libc::c_void,
                        buffer.len() as libc::size_t,
                    )
                } {
                    len if len > 0 => len,
                    _ => return Err(io::Error::last_os_error()),
                };
                let mut ptr = buffer.as_mut_ptr();
                let end = unsafe { buffer.as_ptr().offset(buflen as isize) };
                while (ptr as *const u8) < end {
                    unsafe {
                        let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                        let start =
                            ptr as isize + (*packet).bh_hdrlen as isize - buffer.as_ptr() as isize;
                        self.packets.push_back((
                            start as usize + header_size,
                            (*packet).bh_caplen as usize - header_size,
                        ));
                        let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                        ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                    }
                }
            }
        }
        let (start, mut len) = self.packets.pop_front().unwrap();
        len += buffer_offset;
        // Zero out part that will become fake ethernet header if on loopback.
        for i in (&mut self.read_buffer[start..start + buffer_offset]).iter_mut() {
            *i = 0;
        }
        Ok(&self.read_buffer[start..start + len])
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    #[path = "unix_interfaces.rs"]
    mod interfaces;
    interfaces::interfaces()
}
<|endoftext|>


<commit_before>extern crate yaml_rust;
extern crate regex;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Requst", r"^\x10\x01\x00.\x00\x00")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ascii");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = res.unwrap();
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);
          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut start = 0;
                  for i in 0..buf.len() {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        printables.push(&buf[start..i]);
                        found = false;
                      }
                    }
                  }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ascii", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<commit_msg>Also search for printable multibyte (wide) strings<commit_after>extern crate yaml_rust;
extern crate regex;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Requst", r"^\x10\x01\x00.\x00\x00")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ascii");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = res.unwrap();
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);
          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut mbfound = false;
                  let mut mbstring = String::new();
                  let mut start = 0;
                  for i in 0..c {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        if (i-start == 1) {
                          if (start > 0) && (buf[start-1] == 0) {
                            mbstring.push(buf[i-1] as char);
                            if !mbfound { mbfound = true; }
                          }
                        }
                        else { printables.push(&buf[start..i]); }
                        found = false;
                      }
                      else if mbfound {
                        mbstring.push('\n');
                        mbfound = false;
                      }
                    }
                  }
                  if found { printables.push(&buf[start..c]); }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ascii", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    for line in mbstring.lines() {
                      if line.len() > 3 { println!("% {}", line); }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<|endoftext|>


<commit_before>extern crate discord;
extern crate regex;

use std::collections::{HashSet, HashMap};
use std::env;
use std::cmp::{PartialEq, Eq};
use std::hash::{Hash, Hasher};
use std::time::{Duration, Instant};

use discord::{Discord, State};
use discord::model::{Event, ChannelId, Message, UserId};

use regex::Regex;

const TIMEOUT: u64 = 5 * 60; // 5 minutes

/// Wrapper around a `Regex, UserId` tuple that implements `PartialEq`, `Eq` and `Hash` manually,
/// using the `Regex::as_str()` function as the `Regex` object itself cannot be hashed.
struct Grep(Regex, UserId);

impl PartialEq<Grep> for Grep {
    fn eq(&self, other: &Self) -> bool {
        let Grep(ref regex, id) = *self;
        let Grep(ref other_regex, other_id) = *other;

        regex.as_str() == other_regex.as_str() && id == other_id
    }
}

impl Eq for Grep {}

impl Hash for Grep {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let Grep(ref regex, id) = *self;
        regex.as_str().hash(state);
        id.hash(state);
    }
}

fn handle_command(message: &Message, greps: &mut HashSet<Grep>) -> String {
    let content = &message.content;
    let author = &message.author;
    // split the message
    let content = {
        let index = match content.find(' ') {
            Some(index) => index,
            None => return include_str!("usage.md").into(),
        };
        let (mention, content) = content.split_at(index);
        if !mention.starts_with("<@") {
            return include_str!("usage.md").into();
        }
        &content[1..]
    };
    if content == "help" {
        include_str!("help.md").into()
    } else if content == "list" {
        if greps.iter().any(|&Grep(_, id)| id == author.id) {
            greps
                .iter()
                .filter(|&&Grep(_, id)| id == author.id)
                .map(|&Grep(ref regex, _)| regex)
                .fold(String::new(),
                      |string, regex| format!("{}\n{}", string, regex))
        } else {
            "you have no greps".into()
        }
    } else if content.starts_with("add ") {
        content
            .splitn(2, ' ')
            .nth(1)
            .map(|pattern| match Regex::new(pattern) {
                     Ok(regex) => {
                         if greps
                                .iter()
                                .any(|&Grep(ref regex, id)| {
                                         id == author.id && regex.as_str() == pattern
                                     }) {
                             "Regex already exists".into()
                         } else {
                             greps.insert(Grep(regex, author.id));
                             "Regex added".into()
                         }
                     }
                     Err(error) => format!("Invalid regex. {}", error),
                 })
            .unwrap()
    } else if content.starts_with("remove ") {
        content
            .splitn(2, ' ')
            .nth(1)
            .map(|pattern| {
                let mut removals = false;
                greps.retain(|&Grep(ref regex, id)| if id == author.id &&
                                                       regex.as_str() == pattern {
                                 removals = true;
                                 false
                             } else {
                                 true
                             });
                if removals {
                    format!("Regex {} removed", pattern)
                } else {
                    format!("Regex {} was not found", pattern)
                }
            })
            .unwrap()
    } else if content == "syntax" {
        include_str!("syntax.md").into()
    } else {
        include_str!("usage.md").into()
    }
}

fn handle_message(message: &Message,
                  greps: &HashSet<Grep>,
                  timeouts: &mut HashMap<(UserId, ChannelId), Instant>)
                  -> Option<String> {
    let users: HashSet<_> = greps
        .iter()
        .filter(|&&Grep(ref regex, _)| regex.is_match(&message.content))
        .map(|&Grep(_, id)| id)
        .filter(|&id| id != message.author.id)
        .filter(|&id| match timeouts.get(&(id, message.channel_id)) {
                    Some(instant) => instant.elapsed() > Duration::from_secs(TIMEOUT),
                    None => true,
                })
        .collect();
    if !users.is_empty() {
        Some(users
                 .into_iter()
                 .inspect(|&id| { timeouts.insert((id, message.channel_id), Instant::now()); })
                 .fold("Hey!".into(),
                       |string, id| format!("{} {}", string, id.mention())))
    } else {
        None
    }
}

fn main() {
    // state
    let mut greps = HashSet::new();
    let mut timeouts = HashMap::new();
    // api
    let discord = Discord::from_bot_token(&env::var("DISCORD_BOT_TOKEN")
                                               .expect("DISCORD_BOT_TOKEN not set"))
            .expect("Login Failed");
    let (mut connection, event) = match discord.connect() {
        Ok(conn) => conn,
        Err(e) => panic!("Unable to connect to discord API: {}", e),
    };
    let uid = {
        let state = State::new(event);
        state.user().id
    };
    // main loop time
    while let Ok(event) = connection.recv_event() {
        if let Event::MessageCreate(message) = event {
            let response = if message.author.bot {
                None
            } else if message.mentions.iter().any(|user| user.id == uid) {
                Some(handle_command(&message, &mut greps))
            } else {
                handle_message(&message, &greps, &mut timeouts)
            };
            let channel = message.channel_id;
            if let Some(content) = response {
                let _ = discord.send_message(channel, &content, "", false);
            }
        }
    }
}
<commit_msg>added new game messages<commit_after>extern crate discord;
extern crate regex;

use std::collections::{HashSet, HashMap};
use std::env;
use std::cmp::{PartialEq, Eq};
use std::hash::{Hash, Hasher};
use std::time::{Duration, Instant};

use discord::{Discord, State};
use discord::model::{Event, ChannelId, Message, UserId};

use regex::Regex;

const TIMEOUT: u64 = 5 * 60; // 5 minutes

/// Wrapper around a `Regex, UserId` tuple that implements `PartialEq`, `Eq` and `Hash` manually,
/// using the `Regex::as_str()` function as the `Regex` object itself cannot be hashed.
struct Grep(Regex, UserId);

impl PartialEq<Grep> for Grep {
    fn eq(&self, other: &Self) -> bool {
        let Grep(ref regex, id) = *self;
        let Grep(ref other_regex, other_id) = *other;

        regex.as_str() == other_regex.as_str() && id == other_id
    }
}

impl Eq for Grep {}

impl Hash for Grep {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let Grep(ref regex, id) = *self;
        regex.as_str().hash(state);
        id.hash(state);
    }
}

fn handle_command(message: &Message, greps: &mut HashSet<Grep>) -> String {
    let content = &message.content;
    let author = &message.author;
    // split the message
    let content = {
        let index = match content.find(' ') {
            Some(index) => index,
            None => return include_str!("usage.md").into(),
        };
        let (mention, content) = content.split_at(index);
        if !mention.starts_with("<@") {
            return include_str!("usage.md").into();
        }
        &content[1..]
    };
    if content == "help" {
        include_str!("help.md").into()
    } else if content == "list" {
        if greps.iter().any(|&Grep(_, id)| id == author.id) {
            greps
                .iter()
                .filter(|&&Grep(_, id)| id == author.id)
                .map(|&Grep(ref regex, _)| regex)
                .fold(String::new(),
                      |string, regex| format!("{}\n{}", string, regex))
        } else {
            "you have no greps".into()
        }
    } else if content.starts_with("add ") {
        content
            .splitn(2, ' ')
            .nth(1)
            .map(|pattern| match Regex::new(pattern) {
                     Ok(regex) => {
                         if greps
                                .iter()
                                .any(|&Grep(ref regex, id)| {
                                         id == author.id && regex.as_str() == pattern
                                     }) {
                             "Regex already exists".into()
                         } else {
                             greps.insert(Grep(regex, author.id));
                             "Regex added".into()
                         }
                     }
                     Err(error) => format!("Invalid regex. {}", error),
                 })
            .unwrap()
    } else if content.starts_with("remove ") {
        content
            .splitn(2, ' ')
            .nth(1)
            .map(|pattern| {
                let mut removals = false;
                greps.retain(|&Grep(ref regex, id)| if id == author.id &&
                                                       regex.as_str() == pattern {
                                 removals = true;
                                 false
                             } else {
                                 true
                             });
                if removals {
                    format!("Regex {} removed", pattern)
                } else {
                    format!("Regex {} was not found", pattern)
                }
            })
            .unwrap()
    } else if content == "syntax" {
        include_str!("syntax.md").into()
    } else {
        include_str!("usage.md").into()
    }
}

fn handle_message(message: &Message,
                  greps: &HashSet<Grep>,
                  timeouts: &mut HashMap<(UserId, ChannelId), Instant>)
                  -> Option<String> {
    let users: HashSet<_> = greps
        .iter()
        .filter(|&&Grep(ref regex, _)| regex.is_match(&message.content))
        .map(|&Grep(_, id)| id)
        .filter(|&id| id != message.author.id)
        .filter(|&id| match timeouts.get(&(id, message.channel_id)) {
                    Some(instant) => instant.elapsed() > Duration::from_secs(TIMEOUT),
                    None => true,
                })
        .collect();
    if !users.is_empty() {
        Some(users
                 .into_iter()
                 .inspect(|&id| { timeouts.insert((id, message.channel_id), Instant::now()); })
                 .fold("Hey!".into(),
                       |string, id| format!("{} {}", string, id.mention())))
    } else {
        None
    }
}

fn main() {
    // state
    let mut greps = HashSet::new();
    let mut timeouts = HashMap::new();
    // api
    let discord = Discord::from_bot_token(&env::var("DISCORD_BOT_TOKEN")
                                               .expect("DISCORD_BOT_TOKEN not set"))
            .expect("Login Failed");
    let (mut connection, event) = match discord.connect() {
        Ok(conn) => conn,
        Err(e) => panic!("Unable to connect to discord API: {}", e),
    };
    let uid = {
        let state = State::new(event);
        state.user().id
    };
    connection.set_game_name("I grep things for you".to_string());
    // main loop time
    while let Ok(event) = connection.recv_event() {
        if let Event::MessageCreate(message) = event {
            let response = if message.author.bot {
                None
            } else if message.mentions.iter().any(|user| user.id == uid) {
                Some(handle_command(&message, &mut greps))
            } else {
                handle_message(&message, &greps, &mut timeouts)
            };
            let channel = message.channel_id;
            if let Some(content) = response {
                let _ = discord.send_message(channel, &content, "", false);
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate serde;
extern crate serde_json;

#[macro_use]
extern crate pest;
#[macro_use]
extern crate pest_derive;
#[macro_use]
extern crate serde_derive;

mod grammar;

use grammar::{TurtleParser, Rule};
use pest::{Parser, iterators, inputs};
use std::{fs, env};
use std::io::Read;

#[derive(Serialize, Deserialize)]
struct Result {
    rule: String,
    text: Option<String>,
    children: Vec<Result>,
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let mut file = fs::File::open(&args[1]).unwrap();
    let mut contents = String::new();
    file.read_to_string(&mut contents);

    let pairs = TurtleParser::parse_str(Rule::turtle_doc, &contents)
        .unwrap_or_else(|e| panic!("{}", e));
    let mut result = Vec::new();
    create_result_from_pairs(pairs, &mut result);
    println!("{}", serde_json::to_string(&result).unwrap());
}

fn create_result_from_pairs(pairs: iterators::Pairs<grammar::Rule, inputs::StrInput>, result: &mut Vec<Result>) {
    for pair in pairs {
        let rule = format!("{:?}", pair.as_rule());
        let text = String::from(pair.as_str());
        let mut children = Vec::new();
        create_result_from_pairs(pair.into_inner(), &mut children);
        result.push(Result {
            rule,
            text: if children.len() == 0 {
              Some(text)
            } else {
              None
            },
            children,
        });
    }
}
<commit_msg>Run cargo fmt<commit_after>extern crate serde;
extern crate serde_json;

#[macro_use]
extern crate pest;
#[macro_use]
extern crate pest_derive;
#[macro_use]
extern crate serde_derive;

mod grammar;

use grammar::{TurtleParser, Rule};
use pest::{Parser, iterators, inputs};
use std::{fs, env};
use std::io::Read;

#[derive(Serialize, Deserialize)]
struct Result {
    rule: String,
    text: Option<String>,
    children: Vec<Result>,
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let mut file = fs::File::open(&args[1]).unwrap();
    let mut contents = String::new();
    file.read_to_string(&mut contents);

    let pairs = TurtleParser::parse_str(Rule::turtle_doc, &contents)
        .unwrap_or_else(|e| panic!("{}", e));
    let mut result = Vec::new();
    create_result_from_pairs(pairs, &mut result);
    println!("{}", serde_json::to_string(&result).unwrap());
}

fn create_result_from_pairs(
    pairs: iterators::Pairs<grammar::Rule, inputs::StrInput>,
    result: &mut Vec<Result>,
) {
    for pair in pairs {
        let rule = format!("{:?}", pair.as_rule());
        let text = String::from(pair.as_str());
        let mut children = Vec::new();
        create_result_from_pairs(pair.into_inner(), &mut children);
        result.push(Result {
            rule,
            text: if children.len() == 0 {
                Some(text)
            } else {
                None
            },
            children,
        });
    }
}
<|endoftext|>