ugi 0.2.1

use std::net::SocketAddr;

/// A proxy server to route requests through.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Proxy {
    /// HTTP CONNECT proxy.
    Http {
        addr: SocketAddr,
        auth: Option<ProxyAuth>,
    },
    /// SOCKS5 proxy.
    Socks5 {
        addr: SocketAddr,
        auth: Option<ProxyAuth>,
    },
}

/// Username/password credentials for proxy authentication.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ProxyAuth {
    pub username: String,
    pub password: String,
}

impl Proxy {
    /// Create an HTTP CONNECT proxy without authentication.
    pub fn http(addr: impl Into<SocketAddr>) -> Self {
        Self::Http {
            addr: addr.into(),
            auth: None,
        }
    }

    /// Create an HTTP CONNECT proxy with username/password authentication.
    pub fn http_with_auth(
        addr: impl Into<SocketAddr>,
        username: impl Into<String>,
        password: impl Into<String>,
    ) -> Self {
        Self::Http {
            addr: addr.into(),
            auth: Some(ProxyAuth {
                username: username.into(),
                password: password.into(),
            }),
        }
    }

    /// Create a SOCKS5 proxy without authentication.
    pub fn socks5(addr: impl Into<SocketAddr>) -> Self {
        Self::Socks5 {
            addr: addr.into(),
            auth: None,
        }
    }

    /// Create a SOCKS5 proxy with username/password authentication.
    pub fn socks5_with_auth(
        addr: impl Into<SocketAddr>,
        username: impl Into<String>,
        password: impl Into<String>,
    ) -> Self {
        Self::Socks5 {
            addr: addr.into(),
            auth: Some(ProxyAuth {
                username: username.into(),
                password: password.into(),
            }),
        }
    }

    /// Return the proxy's socket address.
    pub fn addr(&self) -> SocketAddr {
        match self {
            Self::Http { addr, .. } => *addr,
            Self::Socks5 { addr, .. } => *addr,
        }
    }

    /// Return the proxy's authentication credentials, if any.
    pub fn auth(&self) -> Option<&ProxyAuth> {
        match self {
            Self::Http { auth, .. } => auth.as_ref(),
            Self::Socks5 { auth, .. } => auth.as_ref(),
        }
    }

    /// Returns `true` if this is an HTTP CONNECT proxy.
    pub fn is_http(&self) -> bool {
        matches!(self, Self::Http { .. })
    }

    /// Returns `true` if this is a SOCKS5 proxy.
    pub fn is_socks5(&self) -> bool {
        matches!(self, Self::Socks5 { .. })
    }
}

pub(crate) fn build_http_connect_request(
    target_host: &str,
    target_port: u16,
    auth: Option<&ProxyAuth>,
) -> String {
    let mut request = format!(
        "CONNECT {target_host}:{target_port} HTTP/1.1\r\nHost: {target_host}:{target_port}\r\n"
    );
    if let Some(auth) = auth {
        let credentials = format!("{}:{}", auth.username, auth.password);
        let encoded = crate::util::encode_base64(credentials.as_bytes());
        request.push_str(&format!("Proxy-Authorization: Basic {encoded}\r\n"));
    }
    request.push_str("\r\n");
    request
}

impl From<String> for Proxy {
    fn from(s: String) -> Self {
        parse_proxy_url(&s).unwrap_or_else(|| panic!("invalid proxy URL: {}", s))
    }
}

impl From<&str> for Proxy {
    fn from(s: &str) -> Self {
        parse_proxy_url(s).unwrap_or_else(|| panic!("invalid proxy URL: {}", s))
    }
}

fn parse_proxy_url(s: &str) -> Option<Proxy> {
    let s = s.trim();

    if s.starts_with("http://") {
        let rest = &s[7..];
        let (addr_str, auth) = parse_proxy_auth(rest)?;
        let addr: SocketAddr = addr_str.parse().ok()?;
        return Some(Proxy::Http { addr, auth });
    }

    if s.starts_with("socks5://") {
        let rest = &s[9..];
        let (addr_str, auth) = parse_proxy_auth(rest)?;
        let addr: SocketAddr = addr_str.parse().ok()?;
        return Some(Proxy::Socks5 { addr, auth });
    }

    if s.starts_with("socks://") {
        let rest = &s[8..];
        let (addr_str, auth) = parse_proxy_auth(rest)?;
        let addr: SocketAddr = addr_str.parse().ok()?;
        return Some(Proxy::Socks5 { addr, auth });
    }

    let addr: SocketAddr = s.parse().ok()?;
    Some(Proxy::Http { addr, auth: None })
}

fn parse_proxy_auth(s: &str) -> Option<(&str, Option<ProxyAuth>)> {
    if let Some((user_info, rest)) = s.split_once('@') {
        let (username, password) = user_info.split_once(':')?;
        let auth = Some(ProxyAuth {
            username: username.to_string(),
            password: password.to_string(),
        });
        Some((rest, auth))
    } else {
        Some((s, None))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_http_proxy() {
        let proxy = Proxy::from("http://127.0.0.1:8080");
        assert!(proxy.is_http());
        assert_eq!(proxy.addr(), "127.0.0.1:8080".parse().unwrap());
    }

    #[test]
    fn parses_http_proxy_with_auth() {
        let proxy = Proxy::from("http://user:pass@127.0.0.1:8080");
        assert!(proxy.is_http());
        let auth = proxy.auth().unwrap();
        assert_eq!(auth.username, "user");
        assert_eq!(auth.password, "pass");
    }

    #[test]
    fn parses_socks5_proxy() {
        let proxy = Proxy::from("socks5://127.0.0.1:1080");
        assert!(proxy.is_socks5());
        assert_eq!(proxy.addr(), "127.0.0.1:1080".parse().unwrap());
    }

    #[test]
    fn parses_socks5_proxy_with_auth() {
        let proxy = Proxy::from("socks5://user:pass@127.0.0.1:1080");
        assert!(proxy.is_socks5());
        let auth = proxy.auth().unwrap();
        assert_eq!(auth.username, "user");
        assert_eq!(auth.password, "pass");
    }

    #[test]
    fn parses_plain_addr_as_http() {
        let proxy = Proxy::from("127.0.0.1:8080");
        assert!(proxy.is_http());
        assert_eq!(proxy.addr(), "127.0.0.1:8080".parse().unwrap());
    }

    #[test]
    fn builds_http_connect_request_with_http11_request_line() {
        let request = build_http_connect_request("example.com", 443, None);
        assert_eq!(
            request,
            "CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n"
        );
    }

    #[test]
    fn builds_http_connect_request_with_basic_auth() {
        let auth = ProxyAuth {
            username: "user".to_owned(),
            password: "pass".to_owned(),
        };
        let request = build_http_connect_request("example.com", 443, Some(&auth));
        assert!(request.starts_with("CONNECT example.com:443 HTTP/1.1\r\n"));
        assert!(request.contains("\r\nHost: example.com:443\r\n"));
        assert!(request.contains("\r\nProxy-Authorization: Basic dXNlcjpwYXNz\r\n"));
        assert!(request.ends_with("\r\n\r\n"));
    }
}