feat: replacement functionality and final cleanups

2021-03-19 17:04:49 +00:00 · 2021-03-19 17:04:49 +00:00 · 36587867c2
commit 36587867c2
parent c3f3321a57
6 changed files with 166 additions and 47 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -1,5 +1,14 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+[[package]]
+name = "aho-corasick"
+version = "0.7.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7404febffaa47dac81aa44dba71523c9d069b1bdc50a77db41195149e17f68e5"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "ansi_term"
 version = "0.11.0"
@ -77,6 +86,19 @@ dependencies = [
 "vec_map",
 ]

+[[package]]
+name = "env_logger"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "17392a012ea30ef05a610aa97dfb49496e71c9f676b27879922ea5bdf60d9d3f"
+dependencies = [
+ "atty",
+ "humantime",
+ "log",
+ "regex",
+ "termcolor",
+]
+
 [[package]]
 name = "futures"
 version = "0.3.13"
@ -184,6 +206,12 @@ version = "1.3.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "615caabe2c3160b313d52ccc905335f4ed5f10881dd63dc5699d47e90be85691"

+[[package]]
+name = "humantime"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
+
 [[package]]
 name = "instant"
 version = "0.1.9"
@ -362,6 +390,23 @@ dependencies = [
 "bitflags",
 ]

+[[package]]
+name = "regex"
+version = "1.4.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "957056ecddbeba1b26965114e191d2e8589ce74db242b6ea25fc4062427a5c19"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.6.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "24d5f089152e60f62d28b835fbff2cd2e8dc0baf1ac13343bef92ab7eed84548"
+
 [[package]]
 name = "ring"
 version = "0.16.20"
@ -470,6 +515,15 @@ dependencies = [
 "unicode-xid",
 ]

+[[package]]
+name = "termcolor"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2dfed899f0eb03f32ee8c6a0aabdb8a7949659e3466561fc0adf54e26d88c5f4"
+dependencies = [
+ "winapi-util",
+]
+
 [[package]]
 name = "textwrap"
 version = "0.11.0"
@ -484,6 +538,7 @@ name = "tlsproxy"
 version = "0.1.0"
 dependencies = [
 "clap",
+ "env_logger",
 "futures",
 "httparse",
 "rustls",
@ -647,6 +702,15 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"

+[[package]]
+name = "winapi-util"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "winapi-x86_64-pc-windows-gnu"
 version = "0.4.0"
--- a/Cargo.toml
+++ b/Cargo.toml
@ -7,10 +7,11 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

 [dependencies]
+env_logger = "0.8.3"
 tokio = { version = "1", features = ["full"] }
 httparse = "1.3.4"
 futures = "0.3.13"
-rustls = { version = "0.19.0", features = ["dangerous_configuration"] }
+rustls = { version = "0.19.0", features = ["dangerous_configuration", "logging"] }
 webpki = "0.21.4"
 webpki-roots = "0.21.0"
 rustls-pemfile = "0.2.0"
--- a/src/command.rs
+++ b/src/command.rs
@ -1,3 +1,4 @@
+use env_logger;
 use rustls_pemfile;
 use std::io::BufReader;
 use std::{fs, str};
@ -7,15 +8,18 @@ use clap::{App, Arg};
 #[derive(Debug, Clone)]
 pub struct Args {
    pub port: u16,
-    pub certs: String,
+    pub chaincert: String,
    pub key: String,
-    pub replaces: Vec<String>,
+    pub cacert: String,
+    pub replaces: Vec<(String, String)>,
    pub verbose: bool,
 }

 pub fn args() -> Args {
    let version = env!("CARGO_PKG_NAME").to_string() + ", version: " + env!("CARGO_PKG_VERSION");

+    //env_logger::Builder::new().parse_filters("trace").init();
+
    let matches = App::new("tlsproxy")
                        .version(&*version)
                        .author("Mahdi Dibaiee <mdibaiee@pm.me>")
@ -26,10 +30,10 @@ pub fn args() -> Args {
                                .value_name("PORT")
                                .help("Listen on PORT [default: 8080]")
                                .takes_value(true))
-                        .arg(Arg::with_name("certs")
-                                .long("certs")
+                        .arg(Arg::with_name("chaincert")
+                                .long("chaincert")
                                .value_name("FILE")
-                                .help("Read server certificates from FILE. This should contain PEM-format certificates in the right order. The first certificate should certify KEYFILE, the last should be a root CA.")
+                                .help("Read server certificate from FILE. This should contain PEM-format certificate in the right order. The first certificate should certify KEYFILE, the last should be a root CA.")
                                .required(true)
                                .takes_value(true))
                        .arg(Arg::with_name("key")
@ -38,10 +42,16 @@ pub fn args() -> Args {
                                .help("Read private key from FILE. This should be a RSA private key or PKCS8-encoded private key, in PEM format.")
                                .required(true)
                                .takes_value(true))
+                        .arg(Arg::with_name("cacert")
+                                .long("cacert")
+                                .value_name("FILE")
+                                .help("Read CA certificate for client from FILE. This should be a CA certificate for the client.")
+                                .required(true)
+                                .takes_value(true))
                        .arg(Arg::with_name("replace")
                                .long("replace")
                                .value_name("PATTERN")
-                                .help("Replace data in outgoing requests according to patterns specified in the s/MATCH/REPLACEMENT format.")
+                                .help("Replace data in outgoing requests according to a pattern specified in the s/MATCH/REPLACEMENT format. Please note MATCH and REPLACEMENT must be of same length!")
                                .multiple(true)
                                .use_delimiter(true)
                                .takes_value(true))
@ -57,18 +67,30 @@ pub fn args() -> Args {
            .unwrap_or("8080")
            .parse::<u16>()
            .unwrap(),
-        certs: matches
-            .value_of("certs")
+        chaincert: matches
+            .value_of("chaincert")
            .map(|a| a.to_owned())
-            .expect("--certs must be specified"),
+            .expect("--chaincert must be specified"),
        key: matches
            .value_of("key")
            .map(|a| a.to_owned())
            .expect("--key must be specified"),
+        cacert: matches
+            .value_of("cacert")
+            .map(|a| a.to_owned())
+            .expect("--cacert must be specified"),
        verbose: matches.is_present("verbose"),
-        replaces: matches
-            .values_of("replace")
-            .map_or(vec![], |a| a.map(|b| b.to_owned()).collect()),
+        replaces: matches.values_of("replace").map_or(vec![], |a| {
+            a.map(|b| {
+                let mut split = b.split("/").skip(1);
+
+                (
+                    split.next().unwrap().to_owned(),
+                    split.next().unwrap().to_owned(),
+                )
+            })
+            .collect()
+        }),
    };
 }

--- a/src/main.rs
+++ b/src/main.rs
@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::sync::{Arc};
 use std::net::SocketAddr;
 use tokio::net::{TcpListener};
 use rustls::{ServerConfig, ClientConfig, NoClientAuth};
@ -17,27 +17,24 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    println!("Listening on {:#?}", addr);

    let mut server_config = ServerConfig::new(NoClientAuth::new());
-    let certs = command::load_certs(&args.certs);
+    let chaincert = command::load_certs(&args.chaincert);
    let privkey = command::load_private_key(&args.key);
    server_config
-        .set_single_cert(certs.clone(), privkey.clone())
+        .set_single_cert(chaincert, privkey)
        .expect("bad certificates/private key");

    let mut client_config = ClientConfig::new();
-    client_config.key_log = Arc::new(rustls::KeyLogFile::new());

-    client_config
-        .dangerous()
-        .set_certificate_verifier(Arc::new(cert::NoCertificateVerification {}));
-
-    client_config
+    let (added, unused) = client_config
        .root_store
-        .add_pem_file(&mut command::read_file(&args.certs)).unwrap();
-            
-    //client_config
-        //.root_store
-        //.add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
+        .add_pem_file(&mut command::read_file(&args.cacert)).unwrap();

+    println!("{} certificates added, {} unused", added, unused);
+    let (added, unused) = client_config
+        .root_store
+        .add_pem_file(&mut command::read_file(&args.chaincert)).unwrap();
+    println!("{} certificates added, {} unused", added, unused);
+            
    let tcp_listener = TcpListener::bind(addr).await?;
    loop {
        let (tcp_stream, _) = tcp_listener.accept().await?;
--- a/src/proxy.rs
+++ b/src/proxy.rs
@ -4,6 +4,7 @@ use std::error::Error;
 use httparse::Request;
 use webpki;
 use std::str;
+use std::io;
 use std::io::{Write, Read};
 use rustls::{ServerConfig, ClientConfig, Session};

@ -37,30 +38,64 @@ pub async fn proxy(mut incoming: TcpStream, client_config: ClientConfig, server_
            let mut outgoing_std = outgoing.into_std()?;
            let mut outgoing_tls = rustls::Stream::new(&mut client_session, &mut outgoing_std);

-            let mut tmp = vec![0; 1024^2];
-            let mut tmp2 = vec![0; 1024^2];
-            //let ciphersuite = incoming_tls.sess.get_negotiated_ciphersuite().unwrap();
-            //println!("{:#?}", ciphersuite);
+            let mut incoming_buf = vec![0; 4086];
+            let mut outgoing_buf = vec![0; 4086];
            loop {
-                println!("incoming");
-                let incoming_read = ops::sync_read(&mut incoming_tls, &mut tmp);
-                println!("incoming {:#?}", str::from_utf8(&tmp)?);
-                match incoming_read {
-                    Ok(n) if n > 0 => ops::sync_write(&mut outgoing_tls, &tmp),
-                    Ok(0) => return Ok(()),
-                    e => e
-                }.unwrap();
+                match incoming_tls.read(&mut incoming_buf) {
+                    Ok(n) => {
+                        incoming_buf.truncate(n);
+                        if args.verbose && n > 0 {
+                            println!("received data of size {:#?} from client: {:#?}", n, str::from_utf8(&incoming_buf)?);
+                        }

-                let outgoing_read = ops::sync_read(&mut outgoing_tls, &mut tmp2);
+                        args.replaces.iter().for_each(|(from, to)| {
+                            let from_buf = from.as_bytes().to_vec();
+                            let to_buf = to.as_bytes().to_vec();
+                            let from_len = from.len();

-                println!("outgoing {:#?}", str::from_utf8(&tmp2)?);
+                            let mut windowed = incoming_buf.windows(from_len).collect::<Vec<&[u8]>>();
+                            let mut indices: Vec<usize> = vec![];

-                match outgoing_read {
-                    Ok(n) if n > 0 => ops::sync_write(&mut incoming_tls, &tmp2),
-                    Ok(0) => return Ok(()),
+                            while let Some(p) = windowed.iter().position(|x| x.to_owned() == from_buf) {
+                                windowed.push(&to_buf);
+                                windowed.swap_remove(p);
+                                indices.push(p);
+                            }

-                    e => e
-                }.unwrap();
+                            indices.iter().for_each(|p| {
+                                let range = p..&(p+from_len);
+                                incoming_buf.splice(range, to_buf.clone());
+                            });
+                            if args.verbose {
+                                println!("replaced {} instances of {} with {}", indices.len(), from, to);
+                            }
+                        });
+                        outgoing_tls.write_all(&incoming_buf)?;
+                        println!("done writing");
+                    },
+
+                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {},
+
+                    Err(e) => return Err(e.into())
+                }
+
+                match outgoing_tls.read(&mut outgoing_buf) {
+                    Ok(n) => {
+                        outgoing_buf.truncate(n);
+                        if args.verbose && n > 0 {
+                            println!("received data of size {:#?} from server: {:#?}", n, str::from_utf8(&outgoing_buf)?);
+                        }
+
+                        incoming_tls.write_all(&outgoing_buf)?;
+                    }
+                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                        continue;
+                    }
+                    Err(e) => {
+                        println!("error reading from stream: {:#?}", e);
+                        return Err(e.into());
+                    }
+                }
            }
        }

--- a/src/proxy/ops.rs
+++ b/src/proxy/ops.rs
@ -54,7 +54,7 @@ pub fn sync_read<T: Read>(stream: &mut T, buf: &mut Vec<u8>) -> Result<usize, Bo
                buf.truncate(n);
                return Ok(n);
            }
-            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                continue;
            }
            Err(e) => {