<commit_before>extern crate serde;
extern crate serde_json;
#[macro_use] extern crate serde_derive;
#[macro_use] extern crate exonum;
extern crate router;
extern crate bodyparser;
extern crate iron;

use exonum::blockchain::{Blockchain, Service, GenesisConfig, Transaction, ApiContext};
use exonum::messages::{RawTransaction, RawMessage, FromRaw, Message};
use exonum::node::{Node, NodeConfig, NodeApiConfig, TransactionSend};
use exonum::storage::{Fork, MemoryDB, MapIndex};
use exonum::crypto::{PublicKey, Hash};
use exonum::encoding::{self, Field};
use exonum::api::{Api, ApiError};
use iron::prelude::*;
use iron::Handler;
use router::Router;

pub const SERVICE_ID: u16 = 1;

struct CurrencyService;

encoding_struct! {
    struct Wallet {
        const SIZE = 48;

        field pub_key:            &PublicKey  [00 => 32]
        field name:               &str        [32 => 40]
        field balance:            u64         [40 => 48]
    }
}

impl Wallet {
    pub fn increase(&mut self, amount: u64) {
        let balance = self.balance() + amount;
        Field::write(&balance, &mut self.raw, 40, 48);
    }

    pub fn decrease(&mut self, amount: u64) {
        let balance = self.balance() - amount;
        Field::write(&balance, &mut self.raw, 40, 48);
    }
}

pub const TX_WALLET_ID: u16 = 1;

message! {
    struct TxCreateWallet {
        const TYPE = SERVICE_ID;
        const ID = TX_WALLET_ID;
        const SIZE = 40;

        field pub_key:     &PublicKey  [00 => 32]
        field name:        &str        [32 => 40]
    }
}

impl TxCreateWallet {
    pub fn execute(&self, fork: &mut Fork) {
        let mut schema = CurrencySchema::new(fork);
        if let None = schema.wallet(self.pub_key()) {
            let wallet = Wallet::new(self.pub_key(), self.name(), 0);
            schema.wallets().put(self.pub_key(), wallet)
        }
    }
}

impl TxIssue {
    pub fn execute(&self, fork: &mut Fork) {
        let mut schema = CurrencySchema::new(fork);
        if let Some(mut wallet) = schema.wallet(self.pub_key()) {
            wallet.increase(self.amount());
            schema.wallets().put(self.pub_key(), wallet)
        }
    }
}

impl TxTransfer {
    pub fn execute(&self, fork: &mut Fork) {
        let mut schema = CurrencySchema::new(fork);
        let sender = schema.wallet(self.from());
        let receiver = schema.wallet(self.to());
        if let (Some(mut sender), Some(mut receiver)) = (sender, receiver) {
            let amount = self.amount();
            if sender.balance() >= amount {
                sender.decrease(amount);
                receiver.increase(amount);
                println!("{:?} => {:?}", sender, receiver);
                let mut wallets = schema.wallets();
                wallets.put(self.from(), sender);
                wallets.put(self.to(), receiver);
            }
        }
    }
}

pub const TX_ISSUE_ID: u16 = 2;

message! {
    struct TxIssue {
        const TYPE = SERVICE_ID;
        const ID = TX_ISSUE_ID;
        const SIZE = 48;

        field pub_key:     &PublicKey  [00 => 32]
        field amount:      u64         [32 => 40]
        field seed:        u64         [40 => 48]
    }
}

pub const TX_TRANSFER_ID: u16 = 3;

message! {
    struct TxTransfer {
        const TYPE = SERVICE_ID;
        const ID = TX_TRANSFER_ID;
        const SIZE = 80;

        field from:        &PublicKey  [00 => 32]
        field to:          &PublicKey  [32 => 64]
        field amount:      u64         [64 => 72]
        field seed:        u64         [72 => 80]
    }
}

#[derive(Serialize, Deserialize)]
struct TransactionResult {
    tx_hash: Hash,
}

#[derive(Clone)]
struct CryptocurrencyApi<T> {
    channel: T,
}

impl<T: TransactionSend + Clone + 'static> Api for CryptocurrencyApi<T> {
    fn wire(&self, router: &mut Router) {
        let self_ = self.clone();
        let transaction = move |req: &mut Request| -> IronResult<Response> {
            match req.get::<bodyparser::Struct<CurrencyTx>>() {
                Ok(Some(transaction)) => {
                    let tx_hash = transaction.hash();
                    self_.channel.send(transaction).map_err(|e| ApiError::Events(e))?;
                    let json = TransactionResult { tx_hash };
                    self_.ok_response(&serde_json::to_value(&json).unwrap())
                }
                Ok(None) => Err(ApiError::IncorrectRequest("Empty request body".into()))?,
                Err(e) => Err(ApiError::IncorrectRequest(Box::new(e)))?,
            }
        };
        let route_post = "/v1/wallets/transaction";
        router.post(&route_post, transaction, "transaction");
    }
}

#[serde(untagged)]
#[derive(Debug, Clone, Serialize, Deserialize)]
enum CurrencyTx {
    Transfer(TxTransfer),
    Issue(TxIssue),
    CreateWallet(TxCreateWallet),
}

impl Message for CurrencyTx {
    fn raw(&self) -> &RawMessage {
        match *self {
            CurrencyTx::Transfer(ref msg) => msg.raw(),
            CurrencyTx::Issue(ref msg) => msg.raw(),
            CurrencyTx::CreateWallet(ref msg) => msg.raw(),
        }
    }
}

impl FromRaw for CurrencyTx {
    fn from_raw(raw: RawMessage) -> Result<Self, encoding::Error> {
        match raw.message_type() {
            TX_TRANSFER_ID => Ok(CurrencyTx::Transfer(TxTransfer::from_raw(raw)?)),
            TX_ISSUE_ID => Ok(CurrencyTx::Issue(TxIssue::from_raw(raw)?)),
            TX_WALLET_ID => Ok(CurrencyTx::CreateWallet(TxCreateWallet::from_raw(raw)?)),
            _ => Err(encoding::Error::IncorrectMessageType { message_type: raw.message_type() }),
        }
    }
}

impl Transaction for CurrencyTx {
    fn verify(&self) -> bool {
        match *self {
            CurrencyTx::CreateWallet(ref msg) => {
                self.verify_signature(msg.pub_key())
            },
            CurrencyTx::Issue(ref msg) => {
                self.verify_signature(msg.pub_key())
            },
            CurrencyTx::Transfer(ref msg) => {
                self.verify_signature(msg.from()) && (*msg.from() != *msg.to())
            },
        }
    }

    fn execute(&self, view: &mut Fork) {
        match *self {
            CurrencyTx::Transfer(ref msg) => msg.execute(view),
            CurrencyTx::Issue(ref msg) => msg.execute(view),
            CurrencyTx::CreateWallet(ref msg) => msg.execute(view),
        }
    }
}

pub struct CurrencySchema<T> {
    view: T,
}

impl<T> CurrencySchema<T> {
    pub fn new(view: T) -> Self {
        CurrencySchema { view }
    }
}

impl<'a> CurrencySchema<&'a mut Fork> {
    /// Returns `MerklePatriciaTable` with wallets.
    pub fn wallets(&mut self) -> MapIndex<&mut Fork, PublicKey, Wallet> {
        MapIndex::new(vec![20], self.view)
    }

    /// Returns wallet for the given public key.
    pub fn wallet(&mut self, pub_key: &PublicKey) -> Option<Wallet> {
        self.wallets().get(pub_key)
    }

}

impl Service for CurrencyService {
    fn service_name(&self) -> &'static str { "cryptocurrency" }

    fn service_id(&self) -> u16 { SERVICE_ID }

    fn tx_from_raw(&self, raw: RawTransaction) -> Result<Box<Transaction>, encoding::Error> {
        CurrencyTx::from_raw(raw).map(|tx| Box::new(tx) as Box<Transaction>)
    }

    fn public_api_handler(&self, ctx: &ApiContext) -> Option<Box<Handler>> {
        let mut router = Router::new();
        let api = CryptocurrencyApi {
            channel: ctx.node_channel().clone(),
        };
        api.wire(&mut router);
        Some(Box::new(router))
    }
}

fn main() {
    exonum::crypto::init();
    exonum::helpers::init_logger().unwrap();

    let db = MemoryDB::new();
    let services: Vec<Box<Service>> = vec![
        Box::new(CurrencyService),
    ];
    let blockchain = Blockchain::new(Box::new(db), services);

    let (public_key, secret_key) = exonum::crypto::gen_keypair();

    let peer = "0.0.0.0:2000".parse().unwrap();
    let api = "0.0.0.0:8000".parse().unwrap();

    let genesis = GenesisConfig::new(vec![public_key].into_iter());

    let api_cfg = NodeApiConfig {
        enable_blockchain_explorer: true,
        public_api_address: Some(api),
        private_api_address: None,
    };

    let node_cfg = NodeConfig {
        listen_address: peer,
        peers: vec![peer],
        public_key,
        secret_key,
        genesis,
        network: Default::default(),
        whitelist: Default::default(),
        api: api_cfg,
    };

    let mut node = Node::new(blockchain, node_cfg);
    node.run().unwrap();
}
<commit_msg>Reuse the schema instance<commit_after>extern crate serde;
extern crate serde_json;
#[macro_use] extern crate serde_derive;
#[macro_use] extern crate exonum;
extern crate router;
extern crate bodyparser;
extern crate iron;

use exonum::blockchain::{Blockchain, Service, GenesisConfig, Transaction, ApiContext};
use exonum::messages::{RawTransaction, RawMessage, FromRaw, Message};
use exonum::node::{Node, NodeConfig, NodeApiConfig, TransactionSend};
use exonum::storage::{Fork, MemoryDB, MapIndex};
use exonum::crypto::{PublicKey, Hash};
use exonum::encoding::{self, Field};
use exonum::api::{Api, ApiError};
use iron::prelude::*;
use iron::Handler;
use router::Router;

pub const SERVICE_ID: u16 = 1;

struct CurrencyService;

encoding_struct! {
    struct Wallet {
        const SIZE = 48;

        field pub_key:            &PublicKey  [00 => 32]
        field name:               &str        [32 => 40]
        field balance:            u64         [40 => 48]
    }
}

impl Wallet {
    pub fn increase(&mut self, amount: u64) {
        let balance = self.balance() + amount;
        Field::write(&balance, &mut self.raw, 40, 48);
    }

    pub fn decrease(&mut self, amount: u64) {
        let balance = self.balance() - amount;
        Field::write(&balance, &mut self.raw, 40, 48);
    }
}

pub const TX_WALLET_ID: u16 = 1;

message! {
    struct TxCreateWallet {
        const TYPE = SERVICE_ID;
        const ID = TX_WALLET_ID;
        const SIZE = 40;

        field pub_key:     &PublicKey  [00 => 32]
        field name:        &str        [32 => 40]
    }
}

impl TxCreateWallet {
    pub fn execute(&self, mut schema: CurrencySchema) {
        if let None = schema.wallet(self.pub_key()) {
            let wallet = Wallet::new(self.pub_key(), self.name(), 0);
            schema.wallets().put(self.pub_key(), wallet)
        }
    }
}

impl TxIssue {
    pub fn execute(&self, mut schema: CurrencySchema) {
        if let Some(mut wallet) = schema.wallet(self.pub_key()) {
            wallet.increase(self.amount());
            schema.wallets().put(self.pub_key(), wallet)
        }
    }
}

impl TxTransfer {
    pub fn execute(&self, mut schema: CurrencySchema) {
        let sender = schema.wallet(self.from());
        let receiver = schema.wallet(self.to());
        if let (Some(mut sender), Some(mut receiver)) = (sender, receiver) {
            let amount = self.amount();
            if sender.balance() >= amount {
                sender.decrease(amount);
                receiver.increase(amount);
                println!("{:?} => {:?}", sender, receiver);
                let mut wallets = schema.wallets();
                wallets.put(self.from(), sender);
                wallets.put(self.to(), receiver);
            }
        }
    }
}

pub const TX_ISSUE_ID: u16 = 2;

message! {
    struct TxIssue {
        const TYPE = SERVICE_ID;
        const ID = TX_ISSUE_ID;
        const SIZE = 48;

        field pub_key:     &PublicKey  [00 => 32]
        field amount:      u64         [32 => 40]
        field seed:        u64         [40 => 48]
    }
}

pub const TX_TRANSFER_ID: u16 = 3;

message! {
    struct TxTransfer {
        const TYPE = SERVICE_ID;
        const ID = TX_TRANSFER_ID;
        const SIZE = 80;

        field from:        &PublicKey  [00 => 32]
        field to:          &PublicKey  [32 => 64]
        field amount:      u64         [64 => 72]
        field seed:        u64         [72 => 80]
    }
}

#[derive(Serialize, Deserialize)]
struct TransactionResult {
    tx_hash: Hash,
}

#[derive(Clone)]
struct CryptocurrencyApi<T> {
    channel: T,
}

impl<T: TransactionSend + Clone + 'static> Api for CryptocurrencyApi<T> {
    fn wire(&self, router: &mut Router) {
        let self_ = self.clone();
        let transaction = move |req: &mut Request| -> IronResult<Response> {
            match req.get::<bodyparser::Struct<CurrencyTx>>() {
                Ok(Some(transaction)) => {
                    let tx_hash = transaction.hash();
                    self_.channel.send(transaction).map_err(|e| ApiError::Events(e))?;
                    let json = TransactionResult { tx_hash };
                    self_.ok_response(&serde_json::to_value(&json).unwrap())
                }
                Ok(None) => Err(ApiError::IncorrectRequest("Empty request body".into()))?,
                Err(e) => Err(ApiError::IncorrectRequest(Box::new(e)))?,
            }
        };
        let route_post = "/v1/wallets/transaction";
        router.post(&route_post, transaction, "transaction");
    }
}

#[serde(untagged)]
#[derive(Debug, Clone, Serialize, Deserialize)]
enum CurrencyTx {
    Transfer(TxTransfer),
    Issue(TxIssue),
    CreateWallet(TxCreateWallet),
}

impl Message for CurrencyTx {
    fn raw(&self) -> &RawMessage {
        match *self {
            CurrencyTx::Transfer(ref msg) => msg.raw(),
            CurrencyTx::Issue(ref msg) => msg.raw(),
            CurrencyTx::CreateWallet(ref msg) => msg.raw(),
        }
    }
}

impl FromRaw for CurrencyTx {
    fn from_raw(raw: RawMessage) -> Result<Self, encoding::Error> {
        match raw.message_type() {
            TX_TRANSFER_ID => Ok(CurrencyTx::Transfer(TxTransfer::from_raw(raw)?)),
            TX_ISSUE_ID => Ok(CurrencyTx::Issue(TxIssue::from_raw(raw)?)),
            TX_WALLET_ID => Ok(CurrencyTx::CreateWallet(TxCreateWallet::from_raw(raw)?)),
            _ => Err(encoding::Error::IncorrectMessageType { message_type: raw.message_type() }),
        }
    }
}

impl Transaction for CurrencyTx {
    fn verify(&self) -> bool {
        match *self {
            CurrencyTx::CreateWallet(ref msg) => {
                self.verify_signature(msg.pub_key())
            },
            CurrencyTx::Issue(ref msg) => {
                self.verify_signature(msg.pub_key())
            },
            CurrencyTx::Transfer(ref msg) => {
                self.verify_signature(msg.from()) && (*msg.from() != *msg.to())
            },
        }
    }

    fn execute(&self, view: &mut Fork) {
        let schema = CurrencySchema { view };
        match *self {
            CurrencyTx::Transfer(ref msg) => msg.execute(schema),
            CurrencyTx::Issue(ref msg) => msg.execute(schema),
            CurrencyTx::CreateWallet(ref msg) => msg.execute(schema),
        }
    }
}

pub struct CurrencySchema<'a> {
    view: &'a mut Fork,
}

impl<'a> CurrencySchema<'a> {
    pub fn wallets(&mut self) -> MapIndex<&mut Fork, PublicKey, Wallet> {
        MapIndex::new(vec![20], self.view)
    }

    pub fn wallet(&mut self, pub_key: &PublicKey) -> Option<Wallet> {
        self.wallets().get(pub_key)
    }
}

impl Service for CurrencyService {
    fn service_name(&self) -> &'static str { "cryptocurrency" }

    fn service_id(&self) -> u16 { SERVICE_ID }

    fn tx_from_raw(&self, raw: RawTransaction) -> Result<Box<Transaction>, encoding::Error> {
        CurrencyTx::from_raw(raw).map(|tx| Box::new(tx) as Box<Transaction>)
    }

    fn public_api_handler(&self, ctx: &ApiContext) -> Option<Box<Handler>> {
        let mut router = Router::new();
        let api = CryptocurrencyApi {
            channel: ctx.node_channel().clone(),
        };
        api.wire(&mut router);
        Some(Box::new(router))
    }
}

fn main() {
    exonum::crypto::init();
    exonum::helpers::init_logger().unwrap();

    let db = MemoryDB::new();
    let services: Vec<Box<Service>> = vec![
        Box::new(CurrencyService),
    ];
    let blockchain = Blockchain::new(Box::new(db), services);

    let (public_key, secret_key) = exonum::crypto::gen_keypair();

    let peer = "0.0.0.0:2000".parse().unwrap();
    let api = "0.0.0.0:8000".parse().unwrap();

    let genesis = GenesisConfig::new(vec![public_key].into_iter());

    let api_cfg = NodeApiConfig {
        enable_blockchain_explorer: true,
        public_api_address: Some(api),
        private_api_address: None,
    };

    let node_cfg = NodeConfig {
        listen_address: peer,
        peers: vec![peer],
        public_key,
        secret_key,
        genesis,
        network: Default::default(),
        whitelist: Default::default(),
        api: api_cfg,
    };

    let mut node = Node::new(blockchain, node_cfg);
    node.run().unwrap();
}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

const BUFSIZ : usize = 8192;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

// Derived from https://doc.rust-lang.org/src/std/up/src/libstd/io/util.rs.html#46-61
pub fn copy_with_flushes<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> ::std::io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = [0; BUFSIZ];
    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == IoErrorKind::Interrupted => continue,
            Err(ref e) if e.kind() == IoErrorKind::WouldBlock => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        writer.flush()?;
        written += len as u64;
    }
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            copy_with_flushes(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        copy_with_flushes(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let mut request = Client::connect(url)?;

    request.headers.set(
        ::websocket::header::WebSocketProtocol(
            vec!["binary".to_string()]
        )
    );

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.3")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
    `ssh user@host -o ProxyHommand "websocat - ws://..."`
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you can use any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<commit_msg>Use clap's crate_version!()<commit_after>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
#[macro_use] // crate_version
extern crate clap;

const BUFSIZ : usize = 8192;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

// Derived from https://doc.rust-lang.org/src/std/up/src/libstd/io/util.rs.html#46-61
pub fn copy_with_flushes<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> ::std::io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = [0; BUFSIZ];
    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == IoErrorKind::Interrupted => continue,
            Err(ref e) if e.kind() == IoErrorKind::WouldBlock => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        writer.flush()?;
        written += len as u64;
    }
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            copy_with_flushes(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        copy_with_flushes(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let mut request = Client::connect(url)?;

    request.headers.set(
        ::websocket::header::WebSocketProtocol(
            vec!["binary".to_string()]
        )
    );

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version(crate_version!())
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
    `ssh user@host -o ProxyHommand "websocat - ws://..."`
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you can use any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<|endoftext|>


<commit_before>use std::process::{Command, exit};
use std::fs::{OpenOptions, File};
use std::io::prelude::*;
use std::io::{BufReader,BufRead};
use std::error::Error;
use std::path::Path;

fn main() {
    let keys:Vec<String> =
        vec!["AEFB411B072836CD48FF0381AE252C284B5DBA5D".to_string(),
             "9E4F11C6A072942A7B3FD3B0B81EB14A09A25EB0".to_string(),
             "35F52A02854DCCAEC9DD5CC410443C7F54B00041".to_string()];
    let keys_clone:Vec<String> = keys.clone();
    //let mut contents:Vec<String> = vec![];
    gpg_import_keys(keys);
    let init_check:bool = gpg_init();
    if !init_check{
        println!("Update has failed, after gpg keys have been imported.");
        println!("Ensure you run update with \"sudo\", or super user privileges");
        exit(1);
    }else {
        println!("GPG Init: Task Completed Successfully!");
    }
    gpg_remove_keys(keys_clone);
    let mut contents:Vec<String> = check_pacman_conf();
    let confold:bool = create_pacman_conf_old();
    create_new_pacman_conf(contents);
    let keyring:bool = archlabs_keyring();
    let populate:bool = key_populate();
    contents = pacman_conf_final();
    create_new_pacman_conf(contents);
    println!("The automated update portion of the process has completed!");
    println!("Please visit the ArchLABs Forums:");
    println!("https://www.tapatalk.com/groups/archlabs\
    /upgrade-guide-for-existing-installations-t432.html");
    println!("To continue with the manual steps in the process...");
}

fn gpg_import_keys(k:Vec<String>) {

    let mut check:i8 = 0;
    for item in k.iter(){
        let keys:String =
            item.to_string();

        let output = Command::new("gpg")
            .arg("--receive-keys")
            .arg(keys.to_string())
            .output()
            .expect("Process Failed to Execute!");
        if output.stderr.is_empty(){
            println!("{}", String::from_utf8(output.stdout).unwrap());
            check +=1 ; //or default of true
        }else {
            println!("{}", String::from_utf8(output.stderr).unwrap());
            check -= 1;
        }
    }
    if check != 3{
        // One or more keys were not imported.
        print!("{} ", (check*-1));
        println!("keys were not imported");
    }else if check == 3 {
        print!("{} ", (check));
        println!("keys were imported successfully.");
    }
}

fn gpg_init() -> bool{
    let check:bool;
    let output = Command::new("pacman-key")
        .arg("--init")
        .output()
        .expect("Process Failed to Execute!");
    if output.stderr.is_empty(){
        print!("Success: ");
        println!("{}", String::from_utf8(output.stdout).unwrap());
        check = true;
    }else {
        print!("Error: ");
        println!("{}", String::from_utf8(output.stderr).unwrap());
        check = false;
    }
    return check;
}

fn gpg_remove_keys(k:Vec<String>) {

    let mut check:i8 = 0;
    for item in k.iter() {
        let keys:String =
            item.to_string();

        let output =
            Command::new("pacman-key")
            .arg("-r")
            .arg(keys.to_string())
            .output()
            .expect("Process Failed to Execute!");

        if output.stderr.is_empty(){
            println!("{}", String::from_utf8(output.stdout).unwrap());
            check +=1 ; //or default of true
        }else {
            println!("{}", String::from_utf8(output.stderr).unwrap());
            check -= 1;
        }
    }
    if check != 3{
        // One or more keys were not imported.
        print!("{} ", (check*-1));
        println!("keys were not removed!");
    }else if check == 3 {
        print!("{} ", (check));
        println!("keys were removed successfully.");
    }
}

fn create_pacman_conf_old() -> bool{
    let check:bool;
    let output = Command::new("mv")
        .arg("/etc/pacman.conf")
        .arg("/etc/pacman.conf.old")
        .output()
        .expect("Process Failed to Execute!");
    if output.stderr.is_empty(){
        println!("{}", String::from_utf8(output.stdout).unwrap());
        println!("Back up of \"pacman.conf\" complete!");
        check = true;
    }else {
        println!("{}", String::from_utf8(output.stderr).unwrap());
        check = false;
    }
    return check;
}
//This function will be called on two separate iterations
fn create_new_pacman_conf(x:Vec<String>) {
    let path = Path::new("/etc/pacman.conf");
    let display = path.display();
    // Open a file in write-only mode, returns `io::Result<File>`
    let _file = match File::create(&path) {
        Err(why) => panic!("couldn't create {}: {}",
                           display,
                           why.description()),
        Ok(_file) => _file,
    };

    let mut f = OpenOptions::new()
        .write(true)
        .append(true)
        .open("/etc/pacman.conf")
        .unwrap();
    for item in x.iter(){
        if let Err(e) =
        writeln!(f, "{}",item.to_string()){
            println!("{}", e);
        }
    }
}

fn check_pacman_conf() ->Vec<String> {

    let mut contents:Vec<String> = vec![];
    let check_str_one:String =
        "#[archlabs_repo]".to_string();
    let check_str_two:String =
        "[archlabs_repo]".to_string();
    let check_str_three:String =
        "#SigLevel = Never".to_string();
    let check_str_four:String =
        "SigLevel = Never".to_string();
    let check_str_five:String =
        "#Server = https://archlabs.github.io/archlabs_repo/$arch".to_string();
    let check_str_six:String =
        "Server = https://archlabs.github.io/archlabs_repo/$arch".to_string();
    let check_str_seven:String =
        "#Server = https://downloads.sourceforge.net/project/\
        archlabs-repo/archlabs_repo/$arch".to_string();
    let check_str_eight:String =
        "Server = https://downloads.sourceforge.net/project/\
        archlabs-repo/archlabs_repo/$arch".to_string();

    let filename = r"/etc/pacman.conf";
    let f =
        BufReader::new(File::open(filename).unwrap());
    for line in f.lines() {
        match line {
            Ok(line) => if line.contains(&check_str_one) {
                let str_replace =
                    line.replace(&check_str_one, &check_str_two);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_three) {
                let str_replace =
                    line.replace(&check_str_three, &check_str_four);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_five) {
                let str_replace =
                    line.replace(&check_str_five, &check_str_six);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_seven) {
                let str_replace =
                    line.replace(&check_str_seven, &check_str_eight);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else {
                contents.push(line.to_string());
                //println!("{}", line.to_string());
            }
            //Error logic is flawed.. need to workout error check
            //to determine if pacman.conf needs to be modified.
            Err(_) => if check_str_two.eq(line.unwrap().as_str()) {
                println!("No change needed");
            } else {
                println!("Something Happened");
            }
        };
    };
    return contents;
}

fn archlabs_keyring() -> bool {
    let check:bool;
    let output =
        Command::new("pacman")
            .arg("-S")
            .arg("archlabs-keyring")
            .output()
            .expect("Process Failed To Execute!");
    if output.stderr.is_empty(){
        println!("{}", String::from_utf8(output.stdout).unwrap());
        check = true;
    }else {
        println!("{}", String::from_utf8(output.stdout).unwrap());
        check = false;
    }
    return check;
}

fn key_populate() -> bool {
    let check:bool;
    let output =
        Command::new("pacman-key")
            .arg("--populate")
            .arg("archlabs")
            .output()
            .expect("Process Failed To Execute!");
    if output.stderr.is_empty(){
        println!("{}", String::from_utf8(output.stdout).unwrap());
        check = true;
    }else {
        println!("{}", String::from_utf8(output.stderr).unwrap());
        check = false;
    }
    return check;
}

fn pacman_conf_final() -> Vec<String>{

    let mut contents:Vec<String> = vec![];
    let check_str_one:String =
        "[archlabs_repo]".to_string();
    let check_str_two:String =
        "SigLevel = Never".to_string();
    let check_str_three:String =
        "SigLevel = Optional TrustAll".to_string();
    let filename = r"/etc/pacman.conf";
    let f =
        BufReader::new(File::open(filename).unwrap());
    for line in f.lines() {
        match line {
            Ok(line) => if line.contains(&check_str_one) {
                let str_replace =
                    line.replace(&check_str_one, &check_str_one);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_two) {
                let str_replace =
                    line.replace(&check_str_two, &check_str_three);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else {
                contents.push(line.to_string());
                //println!("{}", line.to_string());
            }
            //Error logic is flawed.. need to workout error check
            //to determine if pacman.conf needs to be modified.
            Err(_) => if check_str_two.eq(line.unwrap().as_str()) {
                println!("No change needed");
            } else {
                println!("Something Happened");
            }
        };
    };
    return contents;
}<commit_msg>Few Tweaks to code.<commit_after>use std::process::{Command, exit};
use std::fs::{OpenOptions, File};
use std::io::prelude::*;
use std::io::{BufReader,BufRead};
use std::error::Error;
use std::path::Path;

fn main() {
    let keys:Vec<String> =
        vec!["AEFB411B072836CD48FF0381AE252C284B5DBA5D".to_string(),
             "9E4F11C6A072942A7B3FD3B0B81EB14A09A25EB0".to_string(),
             "35F52A02854DCCAEC9DD5CC410443C7F54B00041".to_string()];
    let keys_clone:Vec<String> = keys.clone();
    //let mut contents:Vec<String> = vec![];
    gpg_import_keys(keys);
    let init_check:bool = gpg_init();
    if !init_check{
        println!("Update has failed, after gpg keys have been imported.");
        println!("Ensure you run update with \"sudo\", or super user privileges");
        exit(1);
    }else {
        println!("GPG Init: Task Completed Successfully!");
    }
    gpg_remove_keys(keys_clone);
    let mut contents:Vec<String> = check_pacman_conf();
    create_pacman_conf_old();
    create_new_pacman_conf(contents);
    archlabs_keyring();
    key_populate();
    contents = pacman_conf_final();
    create_new_pacman_conf(contents);

    println!("The automated update portion of the process has completed!");
    println!("Please visit the ArchLABs Forums:");
    println!("https://www.tapatalk.com/groups/archlabs\
    /upgrade-guide-for-existing-installations-t432.html");
    println!("To continue with the manual steps in the process...");
}

fn gpg_import_keys(k:Vec<String>) {

    let mut check:i8 = 0;
    for item in k.iter(){
        let keys:String =
            item.to_string();

        let output = Command::new("gpg")
            .arg("--receive-keys")
            .arg(keys.to_string())
            .output()
            .expect("Process Failed to Execute!");
        if output.stderr.is_empty(){
            println!("{}", String::from_utf8(output.stdout).unwrap());
            check +=1 ; //or default of true
        }else {
            println!("{}", String::from_utf8(output.stderr).unwrap());
            check -= 1;
        }
    }
    if check != 3{
        // One or more keys were not imported.
        print!("{} ", (check*-1));
        println!("keys were not imported");
    }else if check == 3 {
        print!("{} ", (check));
        println!("keys were imported successfully.");
    }
}

fn gpg_init() -> bool{
    let check:bool;
    let output = Command::new("pacman-key")
        .arg("--init")
        .output()
        .expect("Process Failed to Execute!");
    if output.stderr.is_empty(){
        print!("Success: ");
        println!("{}", String::from_utf8(output.stdout).unwrap());
        check = true;
    }else {
        print!("Error: ");
        println!("{}", String::from_utf8(output.stderr).unwrap());
        check = false;
    }
    return check;
}

fn gpg_remove_keys(k:Vec<String>) {

    let mut check:i8 = 0;
    for item in k.iter() {
        let keys:String =
            item.to_string();

        let output =
            Command::new("pacman-key")
            .arg("-r")
            .arg(keys.to_string())
            .output()
            .expect("Process Failed to Execute!");

        if output.stderr.is_empty(){
            println!("{}", String::from_utf8(output.stdout).unwrap());
            check +=1 ; //or default of true
        }else {
            println!("{}", String::from_utf8(output.stderr).unwrap());
            check -= 1;
        }
    }
    if check != 3{
        // One or more keys were not imported.
        print!("{} ", (check*-1));
        println!("keys were not removed!");
    }else if check == 3 {
        print!("{} ", (check));
        println!("keys were removed successfully.");
    }
}

fn create_pacman_conf_old(){
    let output = Command::new("mv")
        .arg("/etc/pacman.conf")
        .arg("/etc/pacman.conf.old")
        .output()
        .expect("Process Failed to Execute!");
    if output.stderr.is_empty(){
        println!("{}", String::from_utf8(output.stdout).unwrap());
        println!("Back up of \"pacman.conf\" complete!");
    }else {
        println!("{}", String::from_utf8(output.stderr).unwrap());
    }
}
//This function will be called on two separate iterations
fn create_new_pacman_conf(x:Vec<String>) {
    let path = Path::new("/etc/pacman.conf");
    let display = path.display();
    // Open a file in write-only mode, returns `io::Result<File>`
    let _file = match File::create(&path) {
        Err(why) => panic!("couldn't create {}: {}",
                           display,
                           why.description()),
        Ok(_file) => _file,
    };

    let mut f = OpenOptions::new()
        .write(true)
        .append(true)
        .open("/etc/pacman.conf")
        .unwrap();
    for item in x.iter(){
        if let Err(e) =
        writeln!(f, "{}",item.to_string()){
            println!("{}", e);
        }
    }
}

fn check_pacman_conf() ->Vec<String> {

    let mut contents:Vec<String> = vec![];
    let check_str_one:String =
        "#[archlabs_repo]".to_string();
    let check_str_two:String =
        "[archlabs_repo]".to_string();
    let check_str_three:String =
        "#SigLevel = Never".to_string();
    let check_str_four:String =
        "SigLevel = Never".to_string();
    let check_str_five:String =
        "#Server = https://archlabs.github.io/archlabs_repo/$arch".to_string();
    let check_str_six:String =
        "Server = https://archlabs.github.io/archlabs_repo/$arch".to_string();
    let check_str_seven:String =
        "#Server = https://downloads.sourceforge.net/project/\
        archlabs-repo/archlabs_repo/$arch".to_string();
    let check_str_eight:String =
        "Server = https://downloads.sourceforge.net/project/\
        archlabs-repo/archlabs_repo/$arch".to_string();

    let filename = r"/etc/pacman.conf";
    let f =
        BufReader::new(File::open(filename).unwrap());
    for line in f.lines() {
        match line {
            Ok(line) => if line.contains(&check_str_one) {
                let str_replace =
                    line.replace(&check_str_one, &check_str_two);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_three) {
                let str_replace =
                    line.replace(&check_str_three, &check_str_four);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_five) {
                let str_replace =
                    line.replace(&check_str_five, &check_str_six);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_seven) {
                let str_replace =
                    line.replace(&check_str_seven, &check_str_eight);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else {
                contents.push(line.to_string());
                //println!("{}", line.to_string());
            }
            //Error logic is flawed.. need to workout error check
            //to determine if pacman.conf needs to be modified.
            Err(_) => if check_str_two.eq(line.unwrap().as_str()) {
                println!("No change needed");
            } else {
                println!("Something Happened");
            }
        };
    };
    return contents;
}

fn archlabs_keyring() {
    let output =
        Command::new("pacman")
            .arg("-S")
            .arg("archlabs-keyring")
            .output()
            .expect("Process Failed To Execute!");
    if output.stderr.is_empty(){
        println!("{}", String::from_utf8(output.stdout).unwrap());

    }else {
        println!("{}", String::from_utf8(output.stdout).unwrap());
    }
}

fn key_populate() {
    let output =
        Command::new("pacman-key")
            .arg("--populate")
            .arg("archlabs")
            .output()
            .expect("Process Failed To Execute!");
    if output.stderr.is_empty(){
        println!("{}", String::from_utf8(output.stdout).unwrap());
    }else {
        println!("{}", String::from_utf8(output.stderr).unwrap());
    }
}

fn pacman_conf_final() -> Vec<String>{

    let mut contents:Vec<String> = vec![];
    let check_str_one:String =
        "[archlabs_repo]".to_string();
    let check_str_two:String =
        "SigLevel = Never".to_string();
    let check_str_three:String =
        "SigLevel = Optional TrustAll".to_string();
    let filename = r"/etc/pacman.conf";
    let f =
        BufReader::new(File::open(filename).unwrap());
    for line in f.lines() {
        match line {
            Ok(line) => if line.contains(&check_str_one) {
                let str_replace =
                    line.replace(&check_str_one, &check_str_one);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else if line.contains(&check_str_two) {
                let str_replace =
                    line.replace(&check_str_two, &check_str_three);
                contents.push(str_replace.to_string());
                //println!("{}", str_replace);
            } else {
                contents.push(line.to_string());
                //println!("{}", line.to_string());
            }
            //Error logic is flawed.. need to workout error check
            //to determine if pacman.conf needs to be modified.
            Err(_) => if check_str_two.eq(line.unwrap().as_str()) {
                println!("No change needed");
            } else {
                println!("Something Happened");
            }
        };
    };
    return contents;
}<|endoftext|>


<commit_before>
// YAML files as iptables configuration sources

// Rust Core
use std::io::Read;
use std::fs::File;
use std::path::Path;
use std::ascii::AsciiExt;
use std::process::exit;

// Crate.io
extern crate clap;
extern crate yaml_rust;

use clap::{Arg, App};
use yaml_rust::{Yaml, YamlLoader};

fn main() {

    let app = App::new("iptables-init")
        .version("1.0.0")
        .about("\nYAML files as iptables configuration sources")
        .arg(Arg::with_name("CONFIG")
             .multiple(true)
             .short("c")
             .long("config")
             .help("yaml file as iptables configuration source")
             .takes_value(true))
        .args_from_usage("-l --license 'Prints License'");

    let matches = app.get_matches();

    if let Some(ref f_paths) = matches.values_of("CONFIG") {
        for f_path in f_paths.iter() {
            read_yaml(f_path);
        }
    }

}

fn read_yaml(f_path: &str) {

    // Create a path for yaml configuration file
    let path = Path::new(&f_path);
    let display = path.display();

    // Open file
    let mut file = match File::open(&path) {
        Err(_)   => {
            println!("Failed to open file \"{}\"", display);
            exit(1);
        },
        Ok(file) => file,
    };

    // Create buffer string and read content
    let mut buffer_str = String::new();
    match file.read_to_string(&mut buffer_str) {
        Err(_)  => {
            println!("Failed to read file \"{}\"", display);
            exit(1);
        },
        Ok(str) => str,
    };

    // Parse yaml file
    let yaml = match YamlLoader::load_from_str(&buffer_str) {
        Err(_)  => {
            println!("Failed to parse yaml data: \"{}\"", display);
            exit(1);
        },
        Ok(yaml) => yaml,
    };

    // Check if yaml file has no data, exits program if no data
    if yaml.is_empty() {
        println!("\"{}\" has no configuration(s)", display);
        exit(1);
    };

    // Validate configuration template format
    match yaml[0].as_hash() {
        Some(doc) => doc,
        None => {
            println!("Configuration template format is invalid: {}", display);
            exit(1);
        }
    };

    // Parse configuration from YAML configuration file
    parse_yaml(&yaml[0]);

}

fn parse_yaml(doc: &Yaml) {
    // Reset all rules
    reset_rules();
    // Parse `filter` rules
    parse_filter(&doc);
}

fn reset_rules() {
    let mut s:String = "iptables -F".to_string();
    s.push_str("\niptables -t nat -F");
    s.push_str("\niptables -t mangle -F");
    s.push_str("\niptables -X");
    println!("{}", s);
}

fn parse_filter(doc: &Yaml) {
    // Set target name
    let target = "filter";
    // Check if filter section exists
    match doc[target].is_badvalue() {
        false => {
            let labels = ["input", "forward", "output"];
            for label in labels.iter() {
                match_filter(doc, target, label);
            }
        },
        _ => {}
    }
}

fn match_filter(doc: &Yaml, target: &str, label: &str) {
    let rule = doc[target][label].as_str();
    match rule {
        Some("DROP") |
        Some("ACCEPT") |
        Some("REJECT") => println!("iptables -P {} {}", label.to_ascii_uppercase(), rule.unwrap()),
        _ => println!("iptables -P {} ACCEPT", label.to_ascii_uppercase())
    }
}

<commit_msg>Added license, improved parsing and output rules<commit_after>
// YAML files as iptables configuration sources

// Rust Core
use std::io::Read;
use std::fs::File;
use std::path::Path;
use std::process::exit;

// Crate.io
extern crate clap;
extern crate yaml_rust;

use clap::{Arg, App};
use yaml_rust::{Yaml, YamlLoader};

fn main() {

    let app = App::new("iptables-init")
        .version("1.0.0")
        .about("\nYAML files as iptables configuration sources")
        .arg(Arg::with_name("CONFIG")
             .multiple(true)
             .short("c")
             .long("config")
             .help("yaml file as iptables configuration source")
             .conflicts_with("license")
             .takes_value(true))
        .args_from_usage("-l --license 'Prints License'");

    let matches = app.get_matches();

    if let Some(ref f_paths) = matches.values_of("CONFIG") {
        for f_path in f_paths.iter() {
            read_yaml(f_path);
        }
    }

    if matches.is_present("license") {
        print_license();
    }

}

fn print_license() {
    let s:String = "
The MIT License (MIT)

Copyright (c) 2015 Samuel

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.
".to_string();
    println!("{}", s);
}

fn read_yaml(f_path: &str) {

    // Create a path for yaml configuration file
    let path = Path::new(&f_path);
    let display = path.display();

    // Open file
    let mut file = match File::open(&path) {
        Err(_)   => {
            println!("Failed to open file \"{}\"", display);
            exit(1);
        },
        Ok(file) => file,
    };

    // Create buffer string and read content
    let mut buffer_str = String::new();
    match file.read_to_string(&mut buffer_str) {
        Err(_)  => {
            println!("Failed to read file \"{}\"", display);
            exit(1);
        },
        Ok(str) => str,
    };

    // Parse yaml file
    let yaml = match YamlLoader::load_from_str(&buffer_str) {
        Err(_)  => {
            println!("Failed to parse yaml data: \"{}\"", display);
            exit(1);
        },
        Ok(yaml) => yaml,
    };

    // Check if yaml file has no data, exits program if no data
    if yaml.is_empty() {
        println!("\"{}\" has no configuration(s)", display);
        exit(1);
    };

    // Validate configuration template format
    match yaml[0].as_hash() {
        Some(doc) => doc,
        None => {
            println!("Configuration template format is invalid: {}", display);
            exit(1);
        }
    };

    // Parse configuration from YAML configuration file
    parse_yaml(&yaml[0]);

}

fn parse_yaml(doc: &Yaml) {
    // Reset all rules
    reset_rules();
    // Read all rules from yaml
    match doc {
        // Parse if template data is a hash object
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                match k.as_str().unwrap() {
                    // Parse `filter` rules
                    "filter" | "FILTER" => parse_filter(k, v),
                    // Parse custom section
                    _ => parse_section(k, v)
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn reset_rules() {
    let mut s:String = "# reset all rules".to_string();
    s.push_str("\niptables -F");
    s.push_str("\niptables -t nat -F");
    s.push_str("\niptables -t mangle -F");
    s.push_str("\niptables -X");
    println!("{}", s);
}

fn parse_section(id: &Yaml, doc: &Yaml) {
    match doc {
        &Yaml::Hash(ref h) => {
            println!("# {} rules", id.as_str().unwrap());
            for (k, v) in h {
                let k = k.as_str().unwrap();
                println!("{:?} - {:?}", k, v);
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn parse_filter(id: &Yaml, doc: &Yaml) {
    // Check if filter section exists
    match doc {
        &Yaml::Hash(ref h) => {
            println!("# {} rules", id.as_str().unwrap());
            for (k, v) in h {
                let k = k.as_str().unwrap();
                let v = v.as_str().unwrap();
                match k {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        match v {
                            "drop" | "reject" | "accept" | "DROP" | "REJECT" | "ACCEPT" => println!("iptables -P {} {}", k, v),
                            _ => {
                                println!("Rules \"{}\" only accept options of \"drop\",\"reject\" or \"accept\"", k);
                                exit(1);
                            }
                        }
                    },
                    _ => println!("iptables -P {}", k)
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}
<|endoftext|>