herolib-mos 0.3.13

use super::error::{NetworkError, Result};
use std::process::Command;

/// Enable IP forwarding in the kernel
pub fn enable_ip_forwarding() -> Result<()> {
    // Enable IPv4 forwarding
    let output = Command::new("sysctl")
        .args(["-w", "net.ipv4.ip_forward=1"])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to enable IPv4 forwarding: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to enable IPv4 forwarding: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    // Enable IPv6 forwarding
    let output = Command::new("sysctl")
        .args(["-w", "net.ipv6.conf.all.forwarding=1"])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to enable IPv6 forwarding: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to enable IPv6 forwarding: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    Ok(())
}


/// Add iptables FORWARD rule to allow traffic within a bridge
/// This is needed when FORWARD policy is DROP (e.g., Docker)
pub fn add_bridge_forward_rule(bridge: &str) -> Result<()> {
    let check = Command::new("iptables")
        .args(["-C", "FORWARD", "-i", bridge, "-o", bridge, "-j", "ACCEPT"])
        .output();
    
    if check.is_err() || !check.unwrap().status.success() {
        let output = Command::new("iptables")
            .args(["-I", "FORWARD", "-i", bridge, "-o", bridge, "-j", "ACCEPT"])
            .output()
            .map_err(|e| NetworkError::Route(format!("Failed to add bridge forward rule: {}", e)))?;

        if !output.status.success() {
            log::warn!(
                "Failed to add bridge-internal forward rule for {}: {}",
                bridge,
                String::from_utf8_lossy(&output.stderr)
            );
        }
    }
    Ok(())
}

/// Add iptables masquerading rule for outbound traffic from bridge
pub fn add_masquerade_rule(bridge: &str, outbound_interface: &str) -> Result<()> {
    add_bridge_forward_rule(bridge)?;

    // add POSTROUTING rule change source address to outbound interface
    let output = Command::new("iptables")
        .args([
            "-t", "nat",
            "-A", "POSTROUTING",
            "-o", outbound_interface,
            "-s", &format!("{}.0/24", get_bridge_network(bridge)?),
            "-j", "MASQUERADE"
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to add masquerade rule: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to add masquerade rule: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    // Add FORWARD rule to allow traffic from bridge to outbound
    let output = Command::new("iptables")
        .args([
            "-A", "FORWARD",
            "-i", bridge,
            "-o", outbound_interface,
            "-j", "ACCEPT"
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to add forward rule: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to add forward rule: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    // Add FORWARD rule for established/related traffic
    let output = Command::new("iptables")
        .args([
            "-A", "FORWARD",
            "-i", outbound_interface,
            "-o", bridge,
            "-m", "state", "--state", "ESTABLISHED,RELATED",
            "-j", "ACCEPT"
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to add forward rule: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to add forward rule: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    Ok(())
}

/// Delete iptables masquerading rule
pub fn delete_masquerade_rule(bridge: &str, outbound_interface: &str) -> Result<()> {
    // Try to get bridge network, but don't fail if it doesn't have IPv4
    let bridge_network = match get_bridge_network(bridge) {
        Ok(network) => network,
        Err(_) => {
            // Bridge likely doesn't have IPv4 (e.g., Mycelium bridge)
            // Just try to delete any existing rules without specifying source
            return Ok(());
        }
    };
    
    // Delete POSTROUTING rule
    let output = Command::new("iptables")
        .args([
            "-t", "nat",
            "-D", "POSTROUTING",
            "-o", outbound_interface,
            "-s", &format!("{}.0/24", bridge_network),
            "-j", "MASQUERADE"
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to delete masquerade rule: {}", e)))?;

    let stderr = String::from_utf8_lossy(&output.stderr);
    if !output.status.success() && !stderr.contains("No chain/target/match") {
        return Err(NetworkError::Route(format!(
            "Failed to delete masquerade rule: {}",
            stderr
        )));
    }

    // Delete FORWARD rules
    let _ = Command::new("iptables")
        .args(["-D", "FORWARD", "-i", bridge, "-j", "ACCEPT"])
        .output();

    let _ = Command::new("iptables")
        .args([
            "-D", "FORWARD",
            "-i", outbound_interface,
            "-o", bridge,
            "-m", "state", "--state", "ESTABLISHED,RELATED",
            "-j", "ACCEPT"
        ])
        .output();

    Ok(())
}

fn get_bridge_network(bridge: &str) -> Result<String> {
    // Get the IP address of the bridge
    let output = Command::new("ip")
        .args(["addr", "show", bridge])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to get bridge IP: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to get bridge IP: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    let output_str = String::from_utf8_lossy(&output.stdout);
    
    // Parse the IP address from the output
    // Look for "inet <ip>/<mask>" line
    let ip_network = output_str
        .lines()
        .find(|line| line.trim().starts_with("inet "))
        .and_then(|line| {
            line.split_whitespace()
                .nth(1) // Get the second element after "inet"
        })
        .and_then(|addr| {
            // Split on / to get just the network part
            addr.split('/').next()
        });

    match ip_network {
        Some(ip) => {
            // Convert IP to network address by zeroing the host part
            // For IPv4: if IP is 10.200.0.1, return 10.200.0
            let parts: Vec<&str> = ip.split('.').collect();
            if parts.len() == 4 {
                Ok(format!("{}.{}.{}", parts[0], parts[1], parts[2]))
            } else {
                Err(NetworkError::Route(format!(
                    "Invalid IP address format for bridge {}: {}", 
                    bridge, ip
                )))
            }
        }
        None => Err(NetworkError::Route(format!(
            "No IP address found for bridge {}", bridge
        ))),
    }
}

/// Add iptables PREROUTING rule for port forwarding
pub fn add_port_forward_rule(
    external_port: u16,
    dest_ip: &str,
    dest_port: u16,
) -> Result<()> {
    // Add PREROUTING rule to redirect port
    let output = Command::new("iptables")
        .args([
            "-t", "nat",
            "-A", "PREROUTING",
            "-p", "tcp",
            "--dport", &external_port.to_string(),
            "-j", "DNAT",
            "--to-destination", &format!("{}:{}", dest_ip, dest_port)
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to add port forward rule: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to add port forward rule: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    // Add FORWARD rule to allow forwarded traffic
    let output = Command::new("iptables")
        .args([
            "-A", "FORWARD",
            "-p", "tcp",
            "-d", dest_ip,
            "--dport", &dest_port.to_string(),
            "-j", "ACCEPT"
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to add forward rule for port: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to add forward rule for port: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    Ok(())
}

/// Delete iptables PREROUTING rule for port forwarding
pub fn delete_port_forward_rule(
    external_port: u16,
    dest_ip: &str,
    dest_port: u16,
) -> Result<()> {
    // Delete PREROUTING rule
    let output = Command::new("iptables")
        .args([
            "-t", "nat",
            "-D", "PREROUTING",
            "-p", "tcp",
            "--dport", &external_port.to_string(),
            "-j", "DNAT",
            "--to-destination", &format!("{}:{}", dest_ip, dest_port)
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to delete port forward rule: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to delete port forward rule: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    // Delete FORWARD rule
    let output = Command::new("iptables")
        .args([
            "-D", "FORWARD",
            "-p", "tcp",
            "-d", dest_ip,
            "--dport", &dest_port.to_string(),
            "-j", "ACCEPT"
        ])
        .output()
        .map_err(|e| NetworkError::Route(format!("Failed to delete forward rule for port: {}", e)))?;

    if !output.status.success() {
        return Err(NetworkError::Route(format!(
            "Failed to delete forward rule for port: {}",
            String::from_utf8_lossy(&output.stderr)
        )));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_enable_ip_forwarding() {
        // This test requires root privileges to actually enable forwarding
        // To run with root: sudo cargo test test_enable_ip_forwarding -- --ignored

        let result = enable_ip_forwarding();
        assert!(result.is_ok(), "IP forwarding should be enabled");
    }
    
    // Integration tests - require root privileges
    #[cfg(target_os = "linux")]
    #[test]
    #[ignore]
    fn test_add_masquerade_rule() {
        // Test adding masquerade rule
        // This test requires root privileges
        // To run: sudo cargo test test_add_masquerade_rule -- --ignored
        
        // Create bridge first
        use crate::network::bridge;
        let _bridge = bridge::create_bridge("test-br").unwrap();
        
        // Set bridge IP
        let _ = std::process::Command::new("ip")
            .args(["addr", "add", "10.200.0.1/24", "dev", "test-br"])
            .output();
        
        let result = add_masquerade_rule("test-br", "eth0");
        assert!(result.is_ok(), "Masquerade rule addition should succeed");
        
        // Clean up
        let _ = delete_masquerade_rule("test-br", "eth0");
        let _ = bridge::delete_bridge("test-br");
    }
    
    #[cfg(target_os = "linux")]
    #[test]
    #[ignore]
    fn test_delete_masquerade_rule() {
        // Test deleting masquerade rule
        // This test requires root privileges
        // To run: sudo cargo test test_delete_masquerade_rule -- --ignored
        
        // Create bridge first
        use crate::network::bridge;
        let bridge_name = "test-br-del"; // Unique name
        let _bridge = bridge::create_bridge(bridge_name).unwrap();
        
        // Set bridge IP
        let _ = std::process::Command::new("ip")
            .args(["addr", "add", "10.200.0.1/24", "dev", bridge_name])
            .output();
        
        // Add first
        let _ = add_masquerade_rule(bridge_name, "eth0");
        
        // Then delete
        let result = delete_masquerade_rule(bridge_name, "eth0");
        assert!(result.is_ok(), "Masquerade rule deletion should succeed");
        
        // Clean up bridge
        let _ = bridge::delete_bridge(bridge_name);
    }
    
    #[cfg(target_os = "linux")]
    #[test]
    #[ignore]
    fn test_add_port_forward_rule() {
        // Test adding port forward rule
        // This test requires root privileges
        // To run: sudo cargo test test_add_port_forward_rule -- --ignored
        
        let result = add_port_forward_rule(8080, "10.200.0.10", 80);
        assert!(result.is_ok(), "Port forward rule addition should succeed");
        
        // Clean up
        let _ = delete_port_forward_rule(8080, "10.200.0.10", 80);
    }
    
    #[cfg(target_os = "linux")]
    #[test]
    #[ignore]
    fn test_delete_port_forward_rule() {
        // Test deleting port forward rule
        // This test requires root privileges
        // To run: sudo cargo test test_delete_port_forward_rule -- --ignored
        
        // Add first
        let _ = add_port_forward_rule(8080, "10.200.0.10", 80);
        
        // Then delete
        let result = delete_port_forward_rule(8080, "10.200.0.10", 80);
        assert!(result.is_ok(), "Port forward rule deletion should succeed");
    }
}