<commit_before>extern crate clap;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate colored_json;
#[cfg(test)]
extern crate escargot;
extern crate serde_json;
extern crate varlink;
extern crate varlink_parser;
extern crate varlink_stdinterfaces;

use clap::{App, Arg, SubCommand};
use colored_json::{ColorMode, ColoredFormatter, Colour, Output, PrettyFormatter, Style, Styler};
use error::{ErrorKind, Result};
use failure::ResultExt;
use proxy::{handle, handle_connect};
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::path::Path;
use std::str;
use varlink::{
    Connection, GetInterfaceDescriptionReply, MethodCall, OrgVarlinkServiceClient,
    OrgVarlinkServiceInterface,
};
use varlink_parser::{Format, FormatColored, Varlink};
use varlink_stdinterfaces::org_varlink_resolver::{VarlinkClient, VarlinkClientInterface};

#[cfg(test)]
mod test;

mod error;
mod proxy;

fn varlink_format(filename: &str, line_len: Option<&str>, should_colorize: bool) -> Result<()> {
    let mut buffer = String::new();
    File::open(Path::new(filename))?.read_to_string(&mut buffer)?;

    let v = Varlink::from_string(&buffer)?;
    match should_colorize {
        true => println!(
            "{}",
            v.interface
                .get_multiline_colored(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
        false => println!(
            "{}",
            v.interface
                .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
    };
    Ok(())
}

fn varlink_info(
    address: Option<&str>,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let connection = match activate {
        Some(activate) => Connection::with_activate(activate)?,
        None => match bridge {
            Some(bridge) => Connection::with_bridge(bridge)?,
            None => {
                let address = address.unwrap();
                if address.rfind(':').is_none() {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(address.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            address.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                } else {
                    Connection::with_address(&address)?
                }
            }
        },
    };
    let mut call = OrgVarlinkServiceClient::new(connection);
    let info = call.get_info()?;
    let bold = Style::new().bold();

    if should_colorize {
        println!("{} {}", bold.paint("Vendor:"), info.vendor);
        println!("{} {}", bold.paint("Product:"), info.product);
        println!("{} {}", bold.paint("Version:"), info.version);
        println!("{} {}", bold.paint("URL:"), info.url);
        println!("{}", bold.paint("Interfaces:"));
    } else {
        println!("Vendor: {}", info.vendor);
        println!("Product: {}", info.product);
        println!("Version: {}", info.version);
        println!("URL: {}", info.url);
        println!("Interfaces:");
    }
    for i in info.interfaces {
        println!("  {}", i)
    }
    println!();
    Ok(())
}

fn varlink_help(
    url: &str,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    line_len: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let address: &str;
    let interface: &str;

    let connection = if let Some(del) = url.rfind('/') {
        address = &url[0..del];
        interface = &url[(del + 1)..];
        Connection::with_address(&address)?
    } else {
        interface = url;
        match activate {
            Some(activate) => Connection::with_activate(activate)?,
            None => match bridge {
                Some(bridge) => Connection::with_bridge(bridge)?,
                None => {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                }
            },
        }
    };

    if interface.find('.') == None {
        Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress))?
    }

    let mut call = OrgVarlinkServiceClient::new(connection);
    match call.get_interface_description(interface.to_string())? {
        GetInterfaceDescriptionReply {
            description: Some(desc),
        } => match should_colorize {
            true => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline_colored(
                        0,
                        line_len.unwrap_or("80").parse::<usize>().unwrap_or(80),
                    )
            ),
            false => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
            ),
        },
        _ => {
            return Err(ErrorKind::NoDescription(format!("No description for {}", url)).into());
        }
    };

    Ok(())
}

fn varlink_call(
    url: &str,
    args: Option<&str>,
    more: bool,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let resolved_address: String;
    let address: &str;
    let interface: &str;
    let method: &str;

    let connection = match activate {
        Some(activate) => {
            method = url;
            Connection::with_activate(activate)?
        }
        None => match bridge {
            Some(bridge) => {
                method = url;
                Connection::with_bridge(bridge)?
            }
            None => {
                if let Some(del) = url.rfind('/') {
                    address = &url[0..del];
                    method = &url[(del + 1)..];
                    if method.find('.') == None {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                } else {
                    if let Some(del) = url.rfind('.') {
                        interface = &url[0..del];
                        method = url;
                        if method.find('.') == None {
                            return Err(
                                varlink::Error::from(varlink::ErrorKind::InvalidAddress).into()
                            );
                        }
                    } else {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => {
                            resolved_address = r.address.clone();
                            resolved_address.as_ref()
                        }
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                }
                Connection::with_address(address).context(ErrorKind::Connection(address.into()))?
            }
        },
    };

    let args = match args {
        Some(args) => {
            serde_json::from_str(args).context(ErrorKind::SerdeJsonDe(args.to_string()))?
        }
        None => serde_json::Value::Null,
    };

    let mut call = MethodCall::<serde_json::Value, serde_json::Value, varlink::Error>::new(
        connection.clone(),
        String::from(method),
        args,
    );

    let color_mode = match should_colorize {
        true => ColorMode::On,
        false => ColorMode::Off,
    };

    let cf = ColoredFormatter::with_styler(
        PrettyFormatter::new(),
        Styler {
            key: Colour::Cyan.normal(),
            string_value: Colour::Purple.normal(),
            integer_value: Colour::Purple.normal(),
            float_value: Colour::Purple.normal(),
            bool_value: Colour::Purple.normal(),
            nil_value: Colour::Purple.normal(),
            string_include_quotation: false,
            ..Default::default()
        },
    );

    if !more {
        let reply = call.call()?;
        println!("{}", cf.to_colored_json(&reply, color_mode)?);
    } else {
        for reply in call.more()? {
            println!("{}", cf.clone().to_colored_json(&reply?, color_mode)?);
        }
    }

    Ok(())
}

fn varlink_bridge(address: Option<&str>) -> Result<()> {
    let stdin = ::std::io::stdin();
    let stdout = ::std::io::stdout();

    let inbuf =
        unsafe { ::std::io::BufReader::new(::std::fs::File::from_raw_fd(stdin.as_raw_fd())) };
    let outw = unsafe { ::std::fs::File::from_raw_fd(stdout.as_raw_fd()) };

    if let Some(address) = address {
        handle_connect(address, inbuf, outw)?;
    } else {
        handle(inbuf, outw)?;
    }
    Ok(())
}

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

fn main() -> Result<()> {
    let mut app = App::new("varlink")
        .version(VERSION)
        /*
                .arg(
                    Arg::with_name("timeout")
                        .short("t")
                        .long("timeout")
                        .value_name("SECONDS")
                        .help("time in seconds to wait for a reply")
                        .takes_value(true),
                )
        */
        .arg(
            Arg::with_name("color")
                .long("color")
                .possible_values(&["on", "off", "auto"])
                .default_value("auto")
                .help("colorize output"),
        ).arg(
            Arg::with_name("resolver")
                .short("R")
                .long("resolver")
                .value_name("ADDRESS")
                .help("address of the resolver")
                .takes_value(true)
                .required(false)
                .default_value("unix:/run/org.varlink.resolver"),
        ).arg(
            Arg::with_name("bridge")
                .short("b")
                .long("bridge")
                .value_name("COMMAND")
                .help("Command to execute and connect to")
                .takes_value(true)
                .required(false),
        ).arg(
            Arg::with_name("activate")
                .short("A")
                .long("activate")
                .value_name("COMMAND")
                .help("Service to socket-activate and connect to")
                .long_help(
                    "Service to socket-activate and connect to. The temporary UNIX socket \
                     address is exported as $VARLINK_ADDRESS.",
                ).takes_value(true)
                .required(false),
        ).subcommand(
            SubCommand::with_name("bridge")
                .version(VERSION)
                .about("Bridge varlink messages from stdio to services on this machine")
                .long_about(
                    "Bridge varlink messages on stdin and stdout to varlink services on this \
                     machine.",
                ).arg(
                    Arg::with_name("connect")
                        .short("C")
                        .long("connect")
                        .value_name("ADDRESS")
                        .help("connect directly to ADDRESS")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("call")
                .version(VERSION)
                .about("Call a method")
                .long_about("Call METHOD on INTERFACE at ADDRESS. ARGUMENTS must be valid JSON.")
                .arg(
                    Arg::with_name("more")
                        .short("m")
                        .long("more")
                        .help("wait for multiple method returns if supported"),
                ).arg(
                    Arg::with_name("METHOD")
                        .value_name("[ADDRESS/]INTERFACE.METHOD")
                        .required(true),
                ).arg(Arg::with_name("ARGUMENTS").required(false)),
        ).subcommand(
            SubCommand::with_name("format")
                .version(VERSION)
                .about("Format a varlink service file")
                .arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ).arg(
                    Arg::with_name("FILE")
                        .required(true)
                        .help("The varlink interface definition file to format"),
                ),
        ).subcommand(
            SubCommand::with_name("info")
                .version(VERSION)
                .about("Print information about a service")
                .long_about("Prints information about the service running at ADDRESS.")
                .arg(Arg::with_name("ADDRESS").required(false)),
        ).subcommand(
            SubCommand::with_name("help")
                .version(VERSION)
                .about("Print interface description or service information")
                .long_about("Prints information about INTERFACE.")
                .arg(
                    Arg::with_name("INTERFACE")
                        .value_name("[ADDRESS/]INTERFACE")
                        .required(true),
                ).arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("resolve")
                .version(VERSION)
                .about("Resolve an interface name to a varlink address")
                .long_about("Resolve INTERFACE to the varlink address that implements it.")
                .arg(Arg::with_name("INTERFACE").required(true)),
        ).subcommand(
            SubCommand::with_name("completions")
                .version(VERSION)
                .about("Generates completion scripts for your shell")
                .arg(
                    Arg::with_name("SHELL")
                        .required(true)
                        .possible_values(&["bash", "fish", "zsh"])
                        .help("The shell to generate the script for"),
                ),
        );
    let matches = app.clone().get_matches();
    let resolver = matches.value_of("resolver").unwrap();
    let bridge = matches.value_of("bridge");
    let activate = matches.value_of("activate");
    let color = matches.value_of("color").unwrap();
    let color_bool = match color {
        "on" => true,
        "off" => false,
        _ => ColorMode::should_colorize(&Output::StdOut),
    };

    match matches.subcommand() {
        ("completions", Some(sub_matches)) => {
            let shell = sub_matches.value_of("SHELL").unwrap();
            app.gen_completions_to("varlink", shell.parse().unwrap(), &mut io::stdout());
        }
        ("format", Some(sub_matches)) => {
            let filename = sub_matches.value_of("FILE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_format(filename, cols, color_bool)?
        }
        ("info", Some(sub_matches)) => {
            let address = sub_matches.value_of("ADDRESS");
            if address.is_none() & &activate.is_none() & &bridge.is_none() {
                app.print_help().context(ErrorKind::Argument)?;
                println!();
                return Err(
                    ErrorKind::Connection("No ADDRESS or activation or bridge".into()).into(),
                );
            }

            varlink_info(address, resolver, activate, bridge, color_bool)?
        }
        ("bridge", Some(sub_matches)) => {
            let address = sub_matches.value_of("connect");
            varlink_bridge(address)?
        }
        ("help", Some(sub_matches)) => {
            let interface = sub_matches.value_of("INTERFACE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_help(interface, resolver, activate, bridge, cols, color_bool)?
        }
        ("call", Some(sub_matches)) => {
            let method = sub_matches.value_of("METHOD").unwrap();
            let args = sub_matches.value_of("ARGUMENTS");
            let more = sub_matches.is_present("more");

            varlink_call(method, args, more, resolver, activate, bridge, color_bool)?
        }
        (_, _) => {
            app.print_help().context(ErrorKind::Argument)?;
            println!();
        }
    }
    Ok(())
}
<commit_msg>varlink-cli/src/main.rs:call() simplify bold output<commit_after>extern crate clap;
extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate colored_json;
#[cfg(test)]
extern crate escargot;
extern crate serde_json;
extern crate varlink;
extern crate varlink_parser;
extern crate varlink_stdinterfaces;

use clap::{App, Arg, SubCommand};
use colored_json::{ColorMode, ColoredFormatter, Colour, Output, PrettyFormatter, Style, Styler};
use error::{ErrorKind, Result};
use failure::ResultExt;
use proxy::{handle, handle_connect};
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::path::Path;
use std::str;
use varlink::{
    Connection, GetInterfaceDescriptionReply, MethodCall, OrgVarlinkServiceClient,
    OrgVarlinkServiceInterface,
};
use varlink_parser::{Format, FormatColored, Varlink};
use varlink_stdinterfaces::org_varlink_resolver::{VarlinkClient, VarlinkClientInterface};

#[cfg(test)]
mod test;

mod error;
mod proxy;

fn varlink_format(filename: &str, line_len: Option<&str>, should_colorize: bool) -> Result<()> {
    let mut buffer = String::new();
    File::open(Path::new(filename))?.read_to_string(&mut buffer)?;

    let v = Varlink::from_string(&buffer)?;
    match should_colorize {
        true => println!(
            "{}",
            v.interface
                .get_multiline_colored(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
        false => println!(
            "{}",
            v.interface
                .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
        ),
    };
    Ok(())
}

fn varlink_info(
    address: Option<&str>,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let connection = match activate {
        Some(activate) => Connection::with_activate(activate)?,
        None => match bridge {
            Some(bridge) => Connection::with_bridge(bridge)?,
            None => {
                let address = address.unwrap();
                if address.rfind(':').is_none() {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(address.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            address.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                } else {
                    Connection::with_address(&address)?
                }
            }
        },
    };

    let mut call = OrgVarlinkServiceClient::new(connection);
    let info = call.get_info()?;

    let bold: fn(w: &str) -> String = if should_colorize {
        |w| Style::new().bold().paint(w).to_string()
    } else {
        |w| w.to_string()
    };

    println!("{} {}", bold("Vendor:"), info.vendor);
    println!("{} {}", bold("Product:"), info.product);
    println!("{} {}", bold("Version:"), info.version);
    println!("{} {}", bold("URL:"), info.url);
    println!("{}", bold("Interfaces:"));

    for i in info.interfaces {
        println!("  {}", i)
    }
    println!();
    Ok(())
}

fn varlink_help(
    url: &str,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    line_len: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let address: &str;
    let interface: &str;

    let connection = if let Some(del) = url.rfind('/') {
        address = &url[0..del];
        interface = &url[(del + 1)..];
        Connection::with_address(&address)?
    } else {
        interface = url;
        match activate {
            Some(activate) => Connection::with_activate(activate)?,
            None => match bridge {
                Some(bridge) => Connection::with_bridge(bridge)?,
                None => {
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    let address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => r.address.clone(),
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                    Connection::with_address(&address)?
                }
            },
        }
    };

    if interface.find('.') == None {
        Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress))?
    }

    let mut call = OrgVarlinkServiceClient::new(connection);
    match call.get_interface_description(interface.to_string())? {
        GetInterfaceDescriptionReply {
            description: Some(desc),
        } => match should_colorize {
            true => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline_colored(
                        0,
                        line_len.unwrap_or("80").parse::<usize>().unwrap_or(80),
                    )
            ),
            false => println!(
                "{}",
                Varlink::from_string(&desc)?
                    .interface
                    .get_multiline(0, line_len.unwrap_or("80").parse::<usize>().unwrap_or(80))
            ),
        },
        _ => {
            return Err(ErrorKind::NoDescription(format!("No description for {}", url)).into());
        }
    };

    Ok(())
}

fn varlink_call(
    url: &str,
    args: Option<&str>,
    more: bool,
    resolver: &str,
    activate: Option<&str>,
    bridge: Option<&str>,
    should_colorize: bool,
) -> Result<()> {
    let resolved_address: String;
    let address: &str;
    let interface: &str;
    let method: &str;

    let connection = match activate {
        Some(activate) => {
            method = url;
            Connection::with_activate(activate)?
        }
        None => match bridge {
            Some(bridge) => {
                method = url;
                Connection::with_bridge(bridge)?
            }
            None => {
                if let Some(del) = url.rfind('/') {
                    address = &url[0..del];
                    method = &url[(del + 1)..];
                    if method.find('.') == None {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                } else {
                    if let Some(del) = url.rfind('.') {
                        interface = &url[0..del];
                        method = url;
                        if method.find('.') == None {
                            return Err(
                                varlink::Error::from(varlink::ErrorKind::InvalidAddress).into()
                            );
                        }
                    } else {
                        return Err(varlink::Error::from(varlink::ErrorKind::InvalidAddress).into());
                    }
                    let conn = Connection::new(resolver)?;
                    let mut resolver = VarlinkClient::new(conn);
                    address = match resolver.resolve(interface.into()).call() {
                        Ok(r) => {
                            resolved_address = r.address.clone();
                            resolved_address.as_ref()
                        }
                        _ => Err(varlink::Error::from(varlink::ErrorKind::InterfaceNotFound(
                            interface.into(),
                        )))?,
                    };
                }
                Connection::with_address(address).context(ErrorKind::Connection(address.into()))?
            }
        },
    };

    let args = match args {
        Some(args) => {
            serde_json::from_str(args).context(ErrorKind::SerdeJsonDe(args.to_string()))?
        }
        None => serde_json::Value::Null,
    };

    let mut call = MethodCall::<serde_json::Value, serde_json::Value, varlink::Error>::new(
        connection.clone(),
        String::from(method),
        args,
    );

    let color_mode = match should_colorize {
        true => ColorMode::On,
        false => ColorMode::Off,
    };

    let cf = ColoredFormatter::with_styler(
        PrettyFormatter::new(),
        Styler {
            key: Colour::Cyan.normal(),
            string_value: Colour::Purple.normal(),
            integer_value: Colour::Purple.normal(),
            float_value: Colour::Purple.normal(),
            bool_value: Colour::Purple.normal(),
            nil_value: Colour::Purple.normal(),
            string_include_quotation: false,
            ..Default::default()
        },
    );

    if !more {
        let reply = call.call()?;
        println!("{}", cf.to_colored_json(&reply, color_mode)?);
    } else {
        for reply in call.more()? {
            println!("{}", cf.clone().to_colored_json(&reply?, color_mode)?);
        }
    }

    Ok(())
}

fn varlink_bridge(address: Option<&str>) -> Result<()> {
    let stdin = ::std::io::stdin();
    let stdout = ::std::io::stdout();

    let inbuf =
        unsafe { ::std::io::BufReader::new(::std::fs::File::from_raw_fd(stdin.as_raw_fd())) };
    let outw = unsafe { ::std::fs::File::from_raw_fd(stdout.as_raw_fd()) };

    if let Some(address) = address {
        handle_connect(address, inbuf, outw)?;
    } else {
        handle(inbuf, outw)?;
    }
    Ok(())
}

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

fn main() -> Result<()> {
    let mut app = App::new("varlink")
        .version(VERSION)
        /*
                .arg(
                    Arg::with_name("timeout")
                        .short("t")
                        .long("timeout")
                        .value_name("SECONDS")
                        .help("time in seconds to wait for a reply")
                        .takes_value(true),
                )
        */
        .arg(
            Arg::with_name("color")
                .long("color")
                .possible_values(&["on", "off", "auto"])
                .default_value("auto")
                .help("colorize output"),
        ).arg(
            Arg::with_name("resolver")
                .short("R")
                .long("resolver")
                .value_name("ADDRESS")
                .help("address of the resolver")
                .takes_value(true)
                .required(false)
                .default_value("unix:/run/org.varlink.resolver"),
        ).arg(
            Arg::with_name("bridge")
                .short("b")
                .long("bridge")
                .value_name("COMMAND")
                .help("Command to execute and connect to")
                .takes_value(true)
                .required(false),
        ).arg(
            Arg::with_name("activate")
                .short("A")
                .long("activate")
                .value_name("COMMAND")
                .help("Service to socket-activate and connect to")
                .long_help(
                    "Service to socket-activate and connect to. The temporary UNIX socket \
                     address is exported as $VARLINK_ADDRESS.",
                ).takes_value(true)
                .required(false),
        ).subcommand(
            SubCommand::with_name("bridge")
                .version(VERSION)
                .about("Bridge varlink messages from stdio to services on this machine")
                .long_about(
                    "Bridge varlink messages on stdin and stdout to varlink services on this \
                     machine.",
                ).arg(
                    Arg::with_name("connect")
                        .short("C")
                        .long("connect")
                        .value_name("ADDRESS")
                        .help("connect directly to ADDRESS")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("call")
                .version(VERSION)
                .about("Call a method")
                .long_about("Call METHOD on INTERFACE at ADDRESS. ARGUMENTS must be valid JSON.")
                .arg(
                    Arg::with_name("more")
                        .short("m")
                        .long("more")
                        .help("wait for multiple method returns if supported"),
                ).arg(
                    Arg::with_name("METHOD")
                        .value_name("[ADDRESS/]INTERFACE.METHOD")
                        .required(true),
                ).arg(Arg::with_name("ARGUMENTS").required(false)),
        ).subcommand(
            SubCommand::with_name("format")
                .version(VERSION)
                .about("Format a varlink service file")
                .arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ).arg(
                    Arg::with_name("FILE")
                        .required(true)
                        .help("The varlink interface definition file to format"),
                ),
        ).subcommand(
            SubCommand::with_name("info")
                .version(VERSION)
                .about("Print information about a service")
                .long_about("Prints information about the service running at ADDRESS.")
                .arg(Arg::with_name("ADDRESS").required(false)),
        ).subcommand(
            SubCommand::with_name("help")
                .version(VERSION)
                .about("Print interface description or service information")
                .long_about("Prints information about INTERFACE.")
                .arg(
                    Arg::with_name("INTERFACE")
                        .value_name("[ADDRESS/]INTERFACE")
                        .required(true),
                ).arg(
                    Arg::with_name("COLUMNS")
                        .short("c")
                        .long("col")
                        .help("maximum width of the output")
                        .required(false)
                        .takes_value(true),
                ),
        ).subcommand(
            SubCommand::with_name("resolve")
                .version(VERSION)
                .about("Resolve an interface name to a varlink address")
                .long_about("Resolve INTERFACE to the varlink address that implements it.")
                .arg(Arg::with_name("INTERFACE").required(true)),
        ).subcommand(
            SubCommand::with_name("completions")
                .version(VERSION)
                .about("Generates completion scripts for your shell")
                .arg(
                    Arg::with_name("SHELL")
                        .required(true)
                        .possible_values(&["bash", "fish", "zsh"])
                        .help("The shell to generate the script for"),
                ),
        );
    let matches = app.clone().get_matches();
    let resolver = matches.value_of("resolver").unwrap();
    let bridge = matches.value_of("bridge");
    let activate = matches.value_of("activate");
    let color = matches.value_of("color").unwrap();
    let color_bool = match color {
        "on" => true,
        "off" => false,
        _ => ColorMode::should_colorize(&Output::StdOut),
    };

    match matches.subcommand() {
        ("completions", Some(sub_matches)) => {
            let shell = sub_matches.value_of("SHELL").unwrap();
            app.gen_completions_to("varlink", shell.parse().unwrap(), &mut io::stdout());
        }
        ("format", Some(sub_matches)) => {
            let filename = sub_matches.value_of("FILE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_format(filename, cols, color_bool)?
        }
        ("info", Some(sub_matches)) => {
            let address = sub_matches.value_of("ADDRESS");
            if address.is_none() & &activate.is_none() & &bridge.is_none() {
                app.print_help().context(ErrorKind::Argument)?;
                println!();
                return Err(
                    ErrorKind::Connection("No ADDRESS or activation or bridge".into()).into(),
                );
            }

            varlink_info(address, resolver, activate, bridge, color_bool)?
        }
        ("bridge", Some(sub_matches)) => {
            let address = sub_matches.value_of("connect");
            varlink_bridge(address)?
        }
        ("help", Some(sub_matches)) => {
            let interface = sub_matches.value_of("INTERFACE").unwrap();
            let cols = sub_matches.value_of("COLUMNS");

            varlink_help(interface, resolver, activate, bridge, cols, color_bool)?
        }
        ("call", Some(sub_matches)) => {
            let method = sub_matches.value_of("METHOD").unwrap();
            let args = sub_matches.value_of("ARGUMENTS");
            let more = sub_matches.is_present("more");

            varlink_call(method, args, more, resolver, activate, bridge, color_bool)?
        }
        (_, _) => {
            app.print_help().context(ErrorKind::Argument)?;
            println!();
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>extern crate clingo;

use std::env;
use clingo::*;

pub struct OnStatementData<'a> {
    atom: ClingoAstTerm,
    builder: &'a mut ClingoProgramBuilder,
}

struct MyAstHandler;
impl<'a> ClingoAstStatementHandler<OnStatementData<'a>> for MyAstHandler {
    // adds atom enable to all rule bodies
    fn on_statement(stm: &ClingoAstStatement, data: &mut OnStatementData) -> bool {
        // pass through all statements that are not rules
        if stm.get_type() != Ok(ClingoAstStatementType::Rule) {
            data.builder
                .add(stm)
                .expect("Failed to add statement to ProgramBuilder.");
            return true;
        }

        // copy the current rule body
        let body = unsafe { stm.rule() }.body();
        let mut extended_body = std::vec::Vec::with_capacity(body.len() + 1);
        for e in body {
            extended_body.push(e.clone());
        }

        // create atom enable
        let lit = ClingoAstLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstLiteralType::Symbolic,
            &data.atom,
        );
        // add atom enable to the rule body
        let y: ClingoAstBodyLiteral = ClingoAstBodyLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstBodyLiteralType::Literal,
            &lit,
        );
        extended_body.push(y);

        // initialize the rule
        let head = unsafe { stm.rule() }.head();
        let rule = ClingoAstRule::new(head, &extended_body);

        // initialize the statement
        let stm2 = ClingoAstStatement::new_rule(stm.location(), &rule);

        // add the rewritten statement to the program
        data.builder
            .add(&stm2)
            .expect("Failed to add statement to ProgramBuilder.");
        true
    }
}

fn print_model(model: &mut ClingoModel) {
    // retrieve the symbols in the model
    let atoms = model
        .symbols(ClingoShowType::Shown as clingo_show_type_bitset_t)
        .expect("Failed to retrieve symbols in the model.");

    print!("Model:");

    for atom in atoms {
        // retrieve and print the symbol's string
        print!(" {}", atom.to_string().unwrap());
    }
    println!();
}

fn solve(ctl: &mut ClingoControl) {
    let solve_mode = ClingoSolveMode::Yield;
    let assumptions = vec![];

    // get a solve handle
    let handle = ctl.solve(solve_mode, &assumptions)
        .expect("Failed retrieving solve handle.");

    // loop over all models
    loop {
        handle.resume().expect("Failed resume on solve handle.");
        match handle.model() {
            // print the model
            Ok(model) => print_model(model),
            // stop if there are no more models
            Err(_) => break,
        }
    }

    // close the solve handle
    handle
        .get()
        .expect("Failed to get result from solve handle.");
    handle.close().expect("Failed to close solve handle.");
}

fn main() {
    // collect clingo options from the command line
    let options = env::args().skip(1).collect();

    let mut ctl = ClingoControl::new(options, 20).expect("Failed creating ClingoControl.");

    let sym = create_id("enable", true).unwrap();
    let sym2 = sym.clone();

    {
        // get the program builder
        let builder = ctl.program_builder().unwrap();

        // begin building a program
        builder.begin().expect("Failed building logic program.");

        // initialize the location
        let location = create_clingo_location(0, 0, 0, 0, "<rewrite>", "<rewrite>");

        // initilize atom to add
        let atom = ClingoAstTerm::new_symbol(location, sym);

        let mut data = OnStatementData {
            atom: atom,
            builder: builder,
        };

        // get the AST of the program
        parse_program_with_event_handler("a :- not b. b :- not a.", &MyAstHandler, &mut data)
            .expect("Failed to parse logic program.");

        // add the external statement: #external enable.
        let ext = ClingoAstExternal::new(atom, &[]);

        let stm =
            ClingoAstStatement::new_external(location, ClingoAstStatementType::External, &ext);
        data.builder
            .add(&stm)
            .expect("Failed to add statement to ProgramBuilder.");

        // finish building a program
        data.builder
            .end()
            .expect("Failed to finish building a program.");
    }
    // ground the base part
    let part = ClingoPart::new("base", &[]);
    let parts = vec![part];
    ctl.ground(&parts)
        .expect("Failed to ground a logic program.");

    // solve with external enable = false
    println!("Solving with enable = false...");
    solve(&mut ctl);

    // solve with external enable = true
    println!("Solving with enable = true...");
    ctl.assign_external(&sym2, ClingoTruthValue::True)
        .expect("Failed to assign #external enable true.");
    solve(&mut ctl);

    // solve with external enable = false
    println!("Solving with enable = false...");
    ctl.assign_external(&sym2, ClingoTruthValue::False)
        .expect("Failed to assign #external enable false.");
    solve(&mut ctl);
}
<commit_msg>rename parse function<commit_after>extern crate clingo;

use std::env;
use clingo::*;

pub struct OnStatementData<'a> {
    atom: ClingoAstTerm,
    builder: &'a mut ClingoProgramBuilder,
}

struct MyAstHandler;
impl<'a> ClingoAstStatementHandler<OnStatementData<'a>> for MyAstHandler {
    // adds atom enable to all rule bodies
    fn on_statement(stm: &ClingoAstStatement, data: &mut OnStatementData) -> bool {
        // pass through all statements that are not rules
        if stm.get_type() != Ok(ClingoAstStatementType::Rule) {
            data.builder
                .add(stm)
                .expect("Failed to add statement to ProgramBuilder.");
            return true;
        }

        // copy the current rule body
        let body = unsafe { stm.rule() }.body();
        let mut extended_body = std::vec::Vec::with_capacity(body.len() + 1);
        for e in body {
            extended_body.push(e.clone());
        }

        // create atom enable
        let lit = ClingoAstLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstLiteralType::Symbolic,
            &data.atom,
        );
        // add atom enable to the rule body
        let y: ClingoAstBodyLiteral = ClingoAstBodyLiteral::new(
            data.atom.location(),
            ClingoAstSign::None,
            ClingoAstBodyLiteralType::Literal,
            &lit,
        );
        extended_body.push(y);

        // initialize the rule
        let head = unsafe { stm.rule() }.head();
        let rule = ClingoAstRule::new(head, &extended_body);

        // initialize the statement
        let stm2 = ClingoAstStatement::new_rule(stm.location(), &rule);

        // add the rewritten statement to the program
        data.builder
            .add(&stm2)
            .expect("Failed to add statement to ProgramBuilder.");
        true
    }
}

fn print_model(model: &mut ClingoModel) {
    // retrieve the symbols in the model
    let atoms = model
        .symbols(ClingoShowType::Shown as clingo_show_type_bitset_t)
        .expect("Failed to retrieve symbols in the model.");

    print!("Model:");

    for atom in atoms {
        // retrieve and print the symbol's string
        print!(" {}", atom.to_string().unwrap());
    }
    println!();
}

fn solve(ctl: &mut ClingoControl) {
    let solve_mode = ClingoSolveMode::Yield;
    let assumptions = vec![];

    // get a solve handle
    let handle = ctl.solve(solve_mode, &assumptions)
        .expect("Failed retrieving solve handle.");

    // loop over all models
    loop {
        handle.resume().expect("Failed resume on solve handle.");
        match handle.model() {
            // print the model
            Ok(model) => print_model(model),
            // stop if there are no more models
            Err(_) => break,
        }
    }

    // close the solve handle
    handle
        .get()
        .expect("Failed to get result from solve handle.");
    handle.close().expect("Failed to close solve handle.");
}

fn main() {
    // collect clingo options from the command line
    let options = env::args().skip(1).collect();

    let mut ctl = ClingoControl::new(options, 20).expect("Failed creating ClingoControl.");

    let sym = create_id("enable", true).unwrap();
    let sym2 = sym.clone();

    {
        // get the program builder
        let builder = ctl.program_builder().unwrap();

        // begin building a program
        builder.begin().expect("Failed building logic program.");

        // initialize the location
        let location = create_clingo_location(0, 0, 0, 0, "<rewrite>", "<rewrite>");

        // initilize atom to add
        let atom = ClingoAstTerm::new_symbol(location, sym);

        let mut data = OnStatementData {
            atom: atom,
            builder: builder,
        };

        // get the AST of the program
        parse_program("a :- not b. b :- not a.", &MyAstHandler, &mut data)
            .expect("Failed to parse logic program.");

        // add the external statement: #external enable.
        let ext = ClingoAstExternal::new(atom, &[]);

        let stm =
            ClingoAstStatement::new_external(location, ClingoAstStatementType::External, &ext);
        data.builder
            .add(&stm)
            .expect("Failed to add statement to ProgramBuilder.");

        // finish building a program
        data.builder
            .end()
            .expect("Failed to finish building a program.");
    }
    // ground the base part
    let part = ClingoPart::new("base", &[]);
    let parts = vec![part];
    ctl.ground(&parts)
        .expect("Failed to ground a logic program.");

    // solve with external enable = false
    println!("Solving with enable = false...");
    solve(&mut ctl);

    // solve with external enable = true
    println!("Solving with enable = true...");
    ctl.assign_external(&sym2, ClingoTruthValue::True)
        .expect("Failed to assign #external enable true.");
    solve(&mut ctl);

    // solve with external enable = false
    println!("Solving with enable = false...");
    ctl.assign_external(&sym2, ClingoTruthValue::False)
        .expect("Failed to assign #external enable false.");
    solve(&mut ctl);
}
<|endoftext|>


<commit_before>//! The SMTP transport sends emails using the SMTP protocol.
//!
//! This SMTP client follows [RFC
//! 5321](https://tools.ietf.org/html/rfc5321), and is designed to efficiently send emails from an
//! application to a relay email server, as it relies as much as possible on the relay server
//! for sanity and RFC compliance checks.
//!
//! It implements the following extensions:
//!
//! * 8BITMIME ([RFC 6152](https://tools.ietf.org/html/rfc6152))
//! * AUTH ([RFC 4954](http://tools.ietf.org/html/rfc4954)) with PLAIN, LOGIN and XOAUTH2 mechanisms
//! * STARTTLS ([RFC 2487](http://tools.ietf.org/html/rfc2487))
//! * SMTPUTF8 ([RFC 6531](http://tools.ietf.org/html/rfc6531))
//!

use crate::smtp::authentication::{
    Credentials, Mechanism, DEFAULT_ENCRYPTED_MECHANISMS, DEFAULT_UNENCRYPTED_MECHANISMS,
};
use crate::smtp::client::net::ClientTlsParameters;
use crate::smtp::client::net::DEFAULT_TLS_MIN_PROTOCOL;
use crate::smtp::client::InnerClient;
use crate::smtp::commands::*;
use crate::smtp::error::{Error, SmtpResult};
use crate::smtp::extension::{ClientId, Extension, MailBodyParameter, MailParameter, ServerInfo};
use crate::{SendableEmail, Transport};
use log::{debug, info};
use native_tls::TlsConnector;
use std::net::{SocketAddr, ToSocketAddrs};
use std::time::Duration;

pub mod authentication;
pub mod client;
pub mod commands;
pub mod error;
pub mod extension;
#[cfg(feature = "connection-pool")]
pub mod r2d2;
pub mod response;
pub mod util;

// Registered port numbers:
// https://www.iana.
// org/assignments/service-names-port-numbers/service-names-port-numbers.xhtml

/// Default smtp port
pub const SMTP_PORT: u16 = 25;
/// Default submission port
pub const SUBMISSION_PORT: u16 = 587;
/// Default submission over TLS port
pub const SUBMISSIONS_PORT: u16 = 465;

/// How to apply TLS to a client connection
#[derive(Clone)]
#[allow(missing_debug_implementations)]
pub enum ClientSecurity {
    /// Insecure connection only (for testing purposes)
    None,
    /// Start with insecure connection and use `STARTTLS` when available
    Opportunistic(ClientTlsParameters),
    /// Start with insecure connection and require `STARTTLS`
    Required(ClientTlsParameters),
    /// Use TLS wrapped connection
    Wrapper(ClientTlsParameters),
}

/// Configures connection reuse behavior
#[derive(Clone, Debug, Copy)]
#[cfg_attr(
    feature = "serde-impls",
    derive(serde_derive::Serialize, serde_derive::Deserialize)
)]
pub enum ConnectionReuseParameters {
    /// Unlimited connection reuse
    ReuseUnlimited,
    /// Maximum number of connection reuse
    ReuseLimited(u16),
    /// Disable connection reuse, close connection after each transaction
    NoReuse,
}

/// Contains client configuration
#[allow(missing_debug_implementations)]
#[derive(Clone)]
pub struct SmtpClient {
    /// Enable connection reuse
    connection_reuse: ConnectionReuseParameters,
    /// Name sent during EHLO
    hello_name: ClientId,
    /// Credentials
    credentials: Option<Credentials>,
    /// Socket we are connecting to
    server_addr: SocketAddr,
    /// TLS security configuration
    security: ClientSecurity,
    /// Enable UTF8 mailboxes in envelope or headers
    smtp_utf8: bool,
    /// Optional enforced authentication mechanism
    authentication_mechanism: Option<Mechanism>,
    /// Define network timeout
    /// It can be changed later for specific needs (like a different timeout for each SMTP command)
    timeout: Option<Duration>,
}

/// Builder for the SMTP `SmtpTransport`
impl SmtpClient {
    /// Creates a new SMTP client
    ///
    /// Defaults are:
    ///
    /// * No connection reuse
    /// * No authentication
    /// * No SMTPUTF8 support
    /// * A 60 seconds timeout for smtp commands
    ///
    /// Consider using [`SmtpClient::new_simple`] instead, if possible.
    pub fn new<A: ToSocketAddrs>(addr: A, security: ClientSecurity) -> Result<SmtpClient, Error> {
        let mut addresses = addr.to_socket_addrs()?;

        match addresses.next() {
            Some(addr) => Ok(SmtpClient {
                server_addr: addr,
                security,
                smtp_utf8: false,
                credentials: None,
                connection_reuse: ConnectionReuseParameters::NoReuse,
                hello_name: ClientId::hostname(),
                authentication_mechanism: None,
                timeout: Some(Duration::new(60, 0)),
            }),
            None => Err(Error::Resolution),
        }
    }

    /// Simple and secure transport, should be used when possible.
    /// Creates an encrypted transport over submissions port, using the provided domain
    /// to validate TLS certificates.
    pub fn new_simple(domain: &str) -> Result<SmtpClient, Error> {
        let mut tls_builder = TlsConnector::builder();
        tls_builder.min_protocol_version(Some(DEFAULT_TLS_MIN_PROTOCOL));

        let tls_parameters =
            ClientTlsParameters::new(domain.to_string(), tls_builder.build().unwrap());

        SmtpClient::new(
            (domain, SUBMISSIONS_PORT),
            ClientSecurity::Wrapper(tls_parameters),
        )
    }

    /// Creates a new local SMTP client to port 25
    pub fn new_unencrypted_localhost() -> Result<SmtpClient, Error> {
        SmtpClient::new(("localhost", SMTP_PORT), ClientSecurity::None)
    }

    /// Enable SMTPUTF8 if the server supports it
    pub fn smtp_utf8(mut self, enabled: bool) -> SmtpClient {
        self.smtp_utf8 = enabled;
        self
    }

    /// Set the name used during EHLO
    pub fn hello_name(mut self, name: ClientId) -> SmtpClient {
        self.hello_name = name;
        self
    }

    /// Enable connection reuse
    pub fn connection_reuse(mut self, parameters: ConnectionReuseParameters) -> SmtpClient {
        self.connection_reuse = parameters;
        self
    }

    /// Set the client credentials
    pub fn credentials<S: Into<Credentials>>(mut self, credentials: S) -> SmtpClient {
        self.credentials = Some(credentials.into());
        self
    }

    /// Set the authentication mechanism to use
    pub fn authentication_mechanism(mut self, mechanism: Mechanism) -> SmtpClient {
        self.authentication_mechanism = Some(mechanism);
        self
    }

    /// Set the timeout duration
    pub fn timeout(mut self, timeout: Option<Duration>) -> SmtpClient {
        self.timeout = timeout;
        self
    }

    /// Build the SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn transport(self) -> SmtpTransport {
        SmtpTransport::new(self)
    }
}

/// Represents the state of a client
#[derive(Debug)]
struct State {
    /// Panic state
    pub panic: bool,
    /// Connection reuse counter
    pub connection_reuse_count: u16,
}

/// Structure that implements the high level SMTP client
#[allow(missing_debug_implementations)]
pub struct SmtpTransport {
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// SmtpTransport variable states
    state: State,
    /// Information about the client
    client_info: SmtpClient,
    /// Low level client
    client: InnerClient,
}

macro_rules! try_smtp (
    ($err: expr, $client: ident) => ({
        match $err {
            Ok(val) => val,
            Err(err) => {
                if !$client.state.panic {
                    $client.state.panic = true;
                    $client.close();
                }
                return Err(From::from(err))
            },
        }
    })
);

impl<'a> SmtpTransport {
    /// Creates a new SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn new(builder: SmtpClient) -> SmtpTransport {
        let client = InnerClient::new();

        SmtpTransport {
            client,
            server_info: None,
            client_info: builder,
            state: State {
                panic: false,
                connection_reuse_count: 0,
            },
        }
    }

    fn connect(&mut self) -> Result<(), Error> {
        // Check if the connection is still available
        if (self.state.connection_reuse_count > 0) && (!self.client.is_connected()) {
            self.close();
        }

        if self.state.connection_reuse_count > 0 {
            info!(
                "connection already established to {}",
                self.client_info.server_addr
            );
            return Ok(());
        }

        self.client.connect(
            &self.client_info.server_addr,
            self.client_info.timeout,
            match self.client_info.security {
                ClientSecurity::Wrapper(ref tls_parameters) => Some(tls_parameters),
                _ => None,
            },
        )?;

        self.client.set_timeout(self.client_info.timeout)?;
        let _response = self.client.read_response()?;

        // Log the connection
        info!("connection established to {}", self.client_info.server_addr);

        self.ehlo()?;

        match (
            &self.client_info.security.clone(),
            self.server_info
                .as_ref()
                .unwrap()
                .supports_feature(Extension::StartTls),
        ) {
            (&ClientSecurity::Required(_), false) => {
                return Err(From::from("Could not encrypt connection, aborting"));
            }
            (&ClientSecurity::Opportunistic(_), false) => (),
            (&ClientSecurity::None, _) => (),
            (&ClientSecurity::Wrapper(_), _) => (),
            (&ClientSecurity::Opportunistic(ref tls_parameters), true)
            | (&ClientSecurity::Required(ref tls_parameters), true) => {
                try_smtp!(self.client.command(StarttlsCommand), self);
                try_smtp!(self.client.upgrade_tls_stream(tls_parameters), self);

                debug!("connection encrypted");

                // Send EHLO again
                self.ehlo()?;
            }
        }

        if self.client_info.credentials.is_some() {
            let mut found = false;

            // Compute accepted mechanism
            let accepted_mechanisms = match self.client_info.authentication_mechanism {
                Some(mechanism) => vec![mechanism],
                None => {
                    if self.client.is_encrypted() {
                        DEFAULT_ENCRYPTED_MECHANISMS.to_vec()
                    } else {
                        DEFAULT_UNENCRYPTED_MECHANISMS.to_vec()
                    }
                }
            };

            for mechanism in accepted_mechanisms {
                if self
                    .server_info
                    .as_ref()
                    .unwrap()
                    .supports_auth_mechanism(mechanism)
                {
                    found = true;
                    try_smtp!(
                        self.client
                            .auth(mechanism, self.client_info.credentials.as_ref().unwrap(),),
                        self
                    );
                    break;
                }
            }

            if !found {
                info!("No supported authentication mechanisms available");
            }
        }
        Ok(())
    }

    /// Gets the EHLO response and updates server information
    fn ehlo(&mut self) -> SmtpResult {
        // Extended Hello
        let ehlo_response = try_smtp!(
            self.client.command(EhloCommand::new(ClientId::new(
                self.client_info.hello_name.to_string()
            ),)),
            self
        );

        self.server_info = Some(try_smtp!(ServerInfo::from_response(&ehlo_response), self));

        // Print server information
        debug!("server {}", self.server_info.as_ref().unwrap());

        Ok(ehlo_response)
    }

    /// Reset the client state
    pub fn close(&mut self) {
        // Close the SMTP transaction if needed
        self.client.close();

        // Reset the client state
        self.server_info = None;
        self.state.panic = false;
        self.state.connection_reuse_count = 0;
    }
}

impl<'a> Transport<'a> for SmtpTransport {
    type Result = SmtpResult;

    /// Sends an email
    #[cfg_attr(
        feature = "cargo-clippy",
        allow(clippy::match_same_arms, clippy::cyclomatic_complexity)
    )]
    fn send(&mut self, email: SendableEmail) -> SmtpResult {
        let message_id = email.message_id().to_string();

        if !self.client.is_connected() {
            self.connect()?;
        }

        // Mail
        let mut mail_options = vec![];

        if self
            .server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::EightBitMime)
        {
            mail_options.push(MailParameter::Body(MailBodyParameter::EightBitMime));
        }

        if self
            .server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::SmtpUtfEight)
            && self.client_info.smtp_utf8
        {
            mail_options.push(MailParameter::SmtpUtfEight);
        }

        try_smtp!(
            self.client.command(MailCommand::new(
                email.envelope().from().cloned(),
                mail_options,
            )),
            self
        );

        // Log the mail command
        info!(
            "{}: from=<{}>",
            message_id,
            match email.envelope().from() {
                Some(address) => address.to_string(),
                None => "".to_string(),
            }
        );

        // Recipient
        for to_address in email.envelope().to() {
            try_smtp!(
                self.client
                    .command(RcptCommand::new(to_address.clone(), vec![])),
                self
            );
            // Log the rcpt command
            info!("{}: to=<{}>", message_id, to_address);
        }

        // Data
        try_smtp!(self.client.command(DataCommand), self);

        // Message content
        let result = self.client.message(Box::new(email.message()));

        if result.is_ok() {
            // Increment the connection reuse counter
            self.state.connection_reuse_count += 1;

            // Log the message
            info!(
                "{}: conn_use={}, status=sent ({})",
                message_id,
                self.state.connection_reuse_count,
                result
                    .as_ref()
                    .ok()
                    .unwrap()
                    .message
                    .iter()
                    .next()
                    .unwrap_or(&"no response".to_string())
            );
        }

        // Test if we can reuse the existing connection
        match self.client_info.connection_reuse {
            ConnectionReuseParameters::ReuseLimited(limit)
                if self.state.connection_reuse_count >= limit =>
            {
                self.close()
            }
            ConnectionReuseParameters::NoReuse => self.close(),
            _ => (),
        }

        result
    }
}
<commit_msg>fix(smtp) Allow forcing of a specific auth (#358)<commit_after>//! The SMTP transport sends emails using the SMTP protocol.
//!
//! This SMTP client follows [RFC
//! 5321](https://tools.ietf.org/html/rfc5321), and is designed to efficiently send emails from an
//! application to a relay email server, as it relies as much as possible on the relay server
//! for sanity and RFC compliance checks.
//!
//! It implements the following extensions:
//!
//! * 8BITMIME ([RFC 6152](https://tools.ietf.org/html/rfc6152))
//! * AUTH ([RFC 4954](http://tools.ietf.org/html/rfc4954)) with PLAIN, LOGIN and XOAUTH2 mechanisms
//! * STARTTLS ([RFC 2487](http://tools.ietf.org/html/rfc2487))
//! * SMTPUTF8 ([RFC 6531](http://tools.ietf.org/html/rfc6531))
//!

use crate::smtp::authentication::{
    Credentials, Mechanism, DEFAULT_ENCRYPTED_MECHANISMS, DEFAULT_UNENCRYPTED_MECHANISMS,
};
use crate::smtp::client::net::ClientTlsParameters;
use crate::smtp::client::net::DEFAULT_TLS_MIN_PROTOCOL;
use crate::smtp::client::InnerClient;
use crate::smtp::commands::*;
use crate::smtp::error::{Error, SmtpResult};
use crate::smtp::extension::{ClientId, Extension, MailBodyParameter, MailParameter, ServerInfo};
use crate::{SendableEmail, Transport};
use log::{debug, info};
use native_tls::TlsConnector;
use std::net::{SocketAddr, ToSocketAddrs};
use std::time::Duration;

pub mod authentication;
pub mod client;
pub mod commands;
pub mod error;
pub mod extension;
#[cfg(feature = "connection-pool")]
pub mod r2d2;
pub mod response;
pub mod util;

// Registered port numbers:
// https://www.iana.
// org/assignments/service-names-port-numbers/service-names-port-numbers.xhtml

/// Default smtp port
pub const SMTP_PORT: u16 = 25;
/// Default submission port
pub const SUBMISSION_PORT: u16 = 587;
/// Default submission over TLS port
pub const SUBMISSIONS_PORT: u16 = 465;

/// How to apply TLS to a client connection
#[derive(Clone)]
#[allow(missing_debug_implementations)]
pub enum ClientSecurity {
    /// Insecure connection only (for testing purposes)
    None,
    /// Start with insecure connection and use `STARTTLS` when available
    Opportunistic(ClientTlsParameters),
    /// Start with insecure connection and require `STARTTLS`
    Required(ClientTlsParameters),
    /// Use TLS wrapped connection
    Wrapper(ClientTlsParameters),
}

/// Configures connection reuse behavior
#[derive(Clone, Debug, Copy)]
#[cfg_attr(
    feature = "serde-impls",
    derive(serde_derive::Serialize, serde_derive::Deserialize)
)]
pub enum ConnectionReuseParameters {
    /// Unlimited connection reuse
    ReuseUnlimited,
    /// Maximum number of connection reuse
    ReuseLimited(u16),
    /// Disable connection reuse, close connection after each transaction
    NoReuse,
}

/// Contains client configuration
#[allow(missing_debug_implementations)]
#[derive(Clone)]
pub struct SmtpClient {
    /// Enable connection reuse
    connection_reuse: ConnectionReuseParameters,
    /// Name sent during EHLO
    hello_name: ClientId,
    /// Credentials
    credentials: Option<Credentials>,
    /// Socket we are connecting to
    server_addr: SocketAddr,
    /// TLS security configuration
    security: ClientSecurity,
    /// Enable UTF8 mailboxes in envelope or headers
    smtp_utf8: bool,
    /// Optional enforced authentication mechanism
    authentication_mechanism: Option<Mechanism>,
    /// Force use of the set authentication mechanism even if server does not report to support it
    force_set_auth: bool,
    /// Define network timeout
    /// It can be changed later for specific needs (like a different timeout for each SMTP command)
    timeout: Option<Duration>,
}

/// Builder for the SMTP `SmtpTransport`
impl SmtpClient {
    /// Creates a new SMTP client
    ///
    /// Defaults are:
    ///
    /// * No connection reuse
    /// * No authentication
    /// * No SMTPUTF8 support
    /// * A 60 seconds timeout for smtp commands
    ///
    /// Consider using [`SmtpClient::new_simple`] instead, if possible.
    pub fn new<A: ToSocketAddrs>(addr: A, security: ClientSecurity) -> Result<SmtpClient, Error> {
        let mut addresses = addr.to_socket_addrs()?;

        match addresses.next() {
            Some(addr) => Ok(SmtpClient {
                server_addr: addr,
                security,
                smtp_utf8: false,
                credentials: None,
                connection_reuse: ConnectionReuseParameters::NoReuse,
                hello_name: ClientId::hostname(),
                authentication_mechanism: None,
                force_set_auth: false,
                timeout: Some(Duration::new(60, 0)),
            }),
            None => Err(Error::Resolution),
        }
    }

    /// Simple and secure transport, should be used when possible.
    /// Creates an encrypted transport over submissions port, using the provided domain
    /// to validate TLS certificates.
    pub fn new_simple(domain: &str) -> Result<SmtpClient, Error> {
        let mut tls_builder = TlsConnector::builder();
        tls_builder.min_protocol_version(Some(DEFAULT_TLS_MIN_PROTOCOL));

        let tls_parameters =
            ClientTlsParameters::new(domain.to_string(), tls_builder.build().unwrap());

        SmtpClient::new(
            (domain, SUBMISSIONS_PORT),
            ClientSecurity::Wrapper(tls_parameters),
        )
    }

    /// Creates a new local SMTP client to port 25
    pub fn new_unencrypted_localhost() -> Result<SmtpClient, Error> {
        SmtpClient::new(("localhost", SMTP_PORT), ClientSecurity::None)
    }

    /// Enable SMTPUTF8 if the server supports it
    pub fn smtp_utf8(mut self, enabled: bool) -> SmtpClient {
        self.smtp_utf8 = enabled;
        self
    }

    /// Set the name used during EHLO
    pub fn hello_name(mut self, name: ClientId) -> SmtpClient {
        self.hello_name = name;
        self
    }

    /// Enable connection reuse
    pub fn connection_reuse(mut self, parameters: ConnectionReuseParameters) -> SmtpClient {
        self.connection_reuse = parameters;
        self
    }

    /// Set the client credentials
    pub fn credentials<S: Into<Credentials>>(mut self, credentials: S) -> SmtpClient {
        self.credentials = Some(credentials.into());
        self
    }

    /// Set the authentication mechanism to use
    pub fn authentication_mechanism(mut self, mechanism: Mechanism) -> SmtpClient {
        self.authentication_mechanism = Some(mechanism);
        self
    }

    /// Set if the set authentication mechanism should be force
    pub fn force_set_auth(mut self, force: bool) -> SmtpClient {
        self.force_set_auth = force;
        self
    }

    /// Set the timeout duration
    pub fn timeout(mut self, timeout: Option<Duration>) -> SmtpClient {
        self.timeout = timeout;
        self
    }

    /// Build the SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn transport(self) -> SmtpTransport {
        SmtpTransport::new(self)
    }
}

/// Represents the state of a client
#[derive(Debug)]
struct State {
    /// Panic state
    pub panic: bool,
    /// Connection reuse counter
    pub connection_reuse_count: u16,
}

/// Structure that implements the high level SMTP client
#[allow(missing_debug_implementations)]
pub struct SmtpTransport {
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// SmtpTransport variable states
    state: State,
    /// Information about the client
    client_info: SmtpClient,
    /// Low level client
    client: InnerClient,
}

macro_rules! try_smtp (
    ($err: expr, $client: ident) => ({
        match $err {
            Ok(val) => val,
            Err(err) => {
                if !$client.state.panic {
                    $client.state.panic = true;
                    $client.close();
                }
                return Err(From::from(err))
            },
        }
    })
);

impl<'a> SmtpTransport {
    /// Creates a new SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn new(builder: SmtpClient) -> SmtpTransport {
        let client = InnerClient::new();

        SmtpTransport {
            client,
            server_info: None,
            client_info: builder,
            state: State {
                panic: false,
                connection_reuse_count: 0,
            },
        }
    }

    fn connect(&mut self) -> Result<(), Error> {
        // Check if the connection is still available
        if (self.state.connection_reuse_count > 0) && (!self.client.is_connected()) {
            self.close();
        }

        if self.state.connection_reuse_count > 0 {
            info!(
                "connection already established to {}",
                self.client_info.server_addr
            );
            return Ok(());
        }

        self.client.connect(
            &self.client_info.server_addr,
            self.client_info.timeout,
            match self.client_info.security {
                ClientSecurity::Wrapper(ref tls_parameters) => Some(tls_parameters),
                _ => None,
            },
        )?;

        self.client.set_timeout(self.client_info.timeout)?;
        let _response = self.client.read_response()?;

        // Log the connection
        info!("connection established to {}", self.client_info.server_addr);

        self.ehlo()?;

        match (
            &self.client_info.security.clone(),
            self.server_info
                .as_ref()
                .unwrap()
                .supports_feature(Extension::StartTls),
        ) {
            (&ClientSecurity::Required(_), false) => {
                return Err(From::from("Could not encrypt connection, aborting"));
            }
            (&ClientSecurity::Opportunistic(_), false) => (),
            (&ClientSecurity::None, _) => (),
            (&ClientSecurity::Wrapper(_), _) => (),
            (&ClientSecurity::Opportunistic(ref tls_parameters), true)
            | (&ClientSecurity::Required(ref tls_parameters), true) => {
                try_smtp!(self.client.command(StarttlsCommand), self);
                try_smtp!(self.client.upgrade_tls_stream(tls_parameters), self);

                debug!("connection encrypted");

                // Send EHLO again
                self.ehlo()?;
            }
        }

        if self.client_info.credentials.is_some() {
            let mut found = false;

            if !self.client_info.force_set_auth {
                // Compute accepted mechanism
                let accepted_mechanisms = match self.client_info.authentication_mechanism {
                    Some(mechanism) => vec![mechanism],
                    None => {
                        if self.client.is_encrypted() {
                            DEFAULT_ENCRYPTED_MECHANISMS.to_vec()
                        } else {
                            DEFAULT_UNENCRYPTED_MECHANISMS.to_vec()
                        }
                    }
                };

                for mechanism in accepted_mechanisms {
                    if self
                        .server_info
                        .as_ref()
                        .unwrap()
                        .supports_auth_mechanism(mechanism)
                    {
                        found = true;
                        try_smtp!(
                        self.client
                            .auth(mechanism, self.client_info.credentials.as_ref().unwrap(),),
                        self
                    );
                        break;
                    }
                }
            } else {
                try_smtp!(
                self.client.auth(self.client_info.authentication_mechanism.expect(
                    "force_set_auth set to true, but no authentication mechanism set"), self.client_info.credentials.as_ref().unwrap(),),
                    self
                );
                found = true;
            }

            if !found {
                info!("No supported authentication mechanisms available");
            }
        }
        Ok(())
    }

    /// Gets the EHLO response and updates server information
    fn ehlo(&mut self) -> SmtpResult {
        // Extended Hello
        let ehlo_response = try_smtp!(
            self.client.command(EhloCommand::new(ClientId::new(
                self.client_info.hello_name.to_string()
            ),)),
            self
        );

        self.server_info = Some(try_smtp!(ServerInfo::from_response(&ehlo_response), self));

        // Print server information
        debug!("server {}", self.server_info.as_ref().unwrap());

        Ok(ehlo_response)
    }

    /// Reset the client state
    pub fn close(&mut self) {
        // Close the SMTP transaction if needed
        self.client.close();

        // Reset the client state
        self.server_info = None;
        self.state.panic = false;
        self.state.connection_reuse_count = 0;
    }
}

impl<'a> Transport<'a> for SmtpTransport {
    type Result = SmtpResult;

    /// Sends an email
    #[cfg_attr(
        feature = "cargo-clippy",
        allow(clippy::match_same_arms, clippy::cyclomatic_complexity)
    )]
    fn send(&mut self, email: SendableEmail) -> SmtpResult {
        let message_id = email.message_id().to_string();

        if !self.client.is_connected() {
            self.connect()?;
        }

        // Mail
        let mut mail_options = vec![];

        if self
            .server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::EightBitMime)
        {
            mail_options.push(MailParameter::Body(MailBodyParameter::EightBitMime));
        }

        if self
            .server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::SmtpUtfEight)
            && self.client_info.smtp_utf8
        {
            mail_options.push(MailParameter::SmtpUtfEight);
        }

        try_smtp!(
            self.client.command(MailCommand::new(
                email.envelope().from().cloned(),
                mail_options,
            )),
            self
        );

        // Log the mail command
        info!(
            "{}: from=<{}>",
            message_id,
            match email.envelope().from() {
                Some(address) => address.to_string(),
                None => "".to_string(),
            }
        );

        // Recipient
        for to_address in email.envelope().to() {
            try_smtp!(
                self.client
                    .command(RcptCommand::new(to_address.clone(), vec![])),
                self
            );
            // Log the rcpt command
            info!("{}: to=<{}>", message_id, to_address);
        }

        // Data
        try_smtp!(self.client.command(DataCommand), self);

        // Message content
        let result = self.client.message(Box::new(email.message()));

        if result.is_ok() {
            // Increment the connection reuse counter
            self.state.connection_reuse_count += 1;

            // Log the message
            info!(
                "{}: conn_use={}, status=sent ({})",
                message_id,
                self.state.connection_reuse_count,
                result
                    .as_ref()
                    .ok()
                    .unwrap()
                    .message
                    .iter()
                    .next()
                    .unwrap_or(&"no response".to_string())
            );
        }

        // Test if we can reuse the existing connection
        match self.client_info.connection_reuse {
            ConnectionReuseParameters::ReuseLimited(limit)
                if self.state.connection_reuse_count >= limit =>
            {
                self.close()
            }
            ConnectionReuseParameters::NoReuse => self.close(),
            _ => (),
        }

        result
    }
}
<|endoftext|>