# Windows Tunnel Client Setup Script
# This script sets up SSH tunnels for SMB file sharing and printing

param(
    [string]$ServerHost = "	92.39.98.22",  # Change this to your server IP/hostname
    [string]$Username = "datexpay",      # Change this to your Linux username
    [string]$SshPort = "5022",
    [string]$SshKeyPath = "$env:USERPROFILE\.ssh\linux_host_key",
    [string]$LocalSmbPort = "4450",
    [string]$LocalPrintPort = "8631",
    [string]$PrinterIP = "192.168.1.80"     # Change to your printer's IP
)

# Function to write colored output
function Write-Status {
    param($Message, $Color = "Green")
    Write-Host "[+] $Message" -ForegroundColor $Color
}

function Write-Warning {
    param($Message)
    Write-Host "[!] $Message" -ForegroundColor Yellow
}

function Write-Error {
    param($Message)
    Write-Host "[ERROR] $Message" -ForegroundColor Red
}

# Check if OpenSSH Client is installed
function Test-OpenSSHClient {
    try {
        $null = Get-Command ssh -ErrorAction Stop
        return $true
    } catch {
        return $false
    }
}

# Install OpenSSH Client if missing
function Install-OpenSSHClient {
    Write-Status "Installing OpenSSH Client..."
    try {
        Add-WindowsCapability -Online -Name OpenSSH.Client~~~~0.0.1.0
        Write-Status "OpenSSH Client installed successfully"
        return $true
    } catch {
        Write-Error "Failed to install OpenSSH Client: $_"
        return $false
    }
}

# Generate SSH key if it doesn't exist
function New-SSHKey {
    if (-not (Test-Path $SshKeyPath)) {
        Write-Status "Generating new SSH key pair..."
        ssh-keygen -t ed25519 -f $SshKeyPath -N "''"
        Write-Status "SSH key generated at: $SshKeyPath"
        Write-Warning "Don't forget to add the public key to the Linux host:"
        Write-Host "Content of $SshKeyPath.pub:" -ForegroundColor Cyan
        Get-Content "$SshKeyPath.pub"
    } else {
        Write-Status "SSH key already exists: $SshKeyPath"
    }
}

# Test SSH connection
function Test-SSHConnection {
    Write-Status "Testing SSH connection to $ServerHost..."
    try {
        $result = ssh -i $SshKeyPath -p $SshPort -o ConnectTimeout=10 -o BatchMode=yes "$Username@$ServerHost" "echo 'Connection successful'"
        if ($result -contains "Connection successful") {
            Write-Status "SSH connection test passed!"
            return $true
        }
    } catch {
        Write-Error "SSH connection failed: $_"
        Write-Warning "Please check:"
        Write-Host "  - Server hostname/IP: $ServerHost" -ForegroundColor Yellow
        Write-Host "  - SSH port: $SshPort" -ForegroundColor Yellow
        Write-Host "  - Username: $Username" -ForegroundColor Yellow
        Write-Host "  - SSH key authentication" -ForegroundColor Yellow
        return $false
    }
}

# Start the tunnel
function Start-Tunnel {
    Write-Status "Starting SSH tunnel for SMB and printing..."
    
    $tunnelCommand = @(
        "-i", $SshKeyPath,
        "-L", "${LocalSmbPort}:localhost:445",
        "-L", "${LocalPrintPort}:${PrinterIP}:631",
        "-p", $SshPort,
        "-N",
        "-o", "ServerAliveInterval=60",
        "-o", "ServerAliveCountMax=3",
        "$Username@$ServerHost"
    ) -join " "
    
    # Create batch file to start tunnel
    $batchContent = @"
@echo off
chcp 65001 > nul
echo ========================================
echo    SMB + Print Tunnel - Keep Open!
echo ========================================
echo Server: $ServerHost
echo SMB Access: \\localhost\$LocalSmbPort
echo Printer Port: localhost:$LocalPrintPort
echo ========================================
echo Press Ctrl+C to stop the tunnel
echo ========================================

ssh $tunnelCommand
"@
    
    $batchContent | Out-File -FilePath "$env:USERPROFILE\start-tunnel.bat" -Encoding ASCII
    Write-Status "Tunnel starter created: $env:USERPROFILE\start-tunnel.bat"
}

# Map network drive function
function Map-NetworkDrive {
    Write-Status "Mapping network drive..."
    
    # First, remove any existing mapping on this port
    try {
        net use "\\localhost\$LocalSmbPort" /delete /y 2>$null
    } catch { }
    
    # Create the mapping
    try {
        net use "\\localhost\$LocalSmbPort" "/user:$Username" "/persistent:yes"
        Write-Status "Network drive mapped to Z: (\\localhost\$LocalSmbPort)"
    } catch {
        Write-Warning "Could not automatically map drive. You can map manually:"
        Write-Host "  net use Z: \\localhost\$LocalSmbPort /user:$Username" -ForegroundColor Cyan
    }
}

# Add network printer function
function Add-NetworkPrinter {
    Write-Status "Setting up network printer..."
    
    $printerName = "Remote-Office-Printer"
    
    # Remove existing printer if it exists
    try {
        Remove-Printer -Name $printerName -ErrorAction SilentlyContinue
    } catch { }
    
    # Add the printer
    try {
        Add-Printer -ConnectionName "http://localhost:$LocalPrintPort/ipp/print" -Name $printerName
        Write-Status "Network printer added: $printerName"
    } catch {
        Write-Warning "Could not auto-add printer. You can add manually:"
        Write-Host "  - Open 'Printers & scanners'" -ForegroundColor Cyan
        Write-Host "  - Add printer -> The printer that I want isn't listed" -ForegroundColor Cyan
        Write-Host "  - Add using TCP/IP address" -ForegroundColor Cyan
        Write-Host "  - Hostname: localhost, Port: $LocalPrintPort" -ForegroundColor Cyan
    }
}

# Create management script
function New-ManagementScript {
    $ps1Content = @'
# Tunnel Management Script
param(
    [string]$Action = "status"
)

switch ($Action.ToLower()) {
    "start" {
        Write-Host "Starting tunnel..." -ForegroundColor Green
        Start-Process "cmd.exe" -ArgumentList "/c `"$env:USERPROFILE\start-tunnel.bat`""
    }
    "status" {
        Write-Host "=== Tunnel Status ===" -ForegroundColor Cyan
        $sshProcess = Get-Process ssh -ErrorAction SilentlyContinue
        if ($sshProcess) {
            Write-Host "SSH Tunnel: RUNNING (PID: $($sshProcess.Id))" -ForegroundColor Green
        } else {
            Write-Host "SSH Tunnel: STOPPED" -ForegroundColor Red
        }
        
        Write-Host "`n=== Network Drives ===" -ForegroundColor Cyan
        net use | findstr "localhost"
        
        Write-Host "`n=== Printers ===" -ForegroundColor Cyan
        Get-Printer | Where-Object {$_.Name -like "*Remote*"} | Format-Table Name, ComputerName, Type
        
        Write-Host "`n=== Network Connections ===" -ForegroundColor Cyan
        netstat -an | findstr ":$using:LocalSmbPort"
        netstat -an | findstr ":$using:LocalPrintPort"
    }
    "stop" {
        Write-Host "Stopping all SSH tunnels..." -ForegroundColor Yellow
        Get-Process ssh -ErrorAction SilentlyContinue | Stop-Process -Force
    }
    "test" {
        Write-Host "Testing connections..." -ForegroundColor Cyan
        Test-NetConnection -ComputerName localhost -Port $using:LocalSmbPort
        Test-NetConnection -ComputerName localhost -Port $using:LocalPrintPort
    }
    default {
        Write-Host "Usage: .\manage-tunnel.ps1 {start|status|stop|test}" -ForegroundColor Yellow
    }
}
'@
    
    $ps1Content | Out-File -FilePath ".\manage-tunnel.ps1" -Encoding UTF8
    Write-Status "Management script created: .\manage-tunnel.ps1"
}

# Main execution
Write-Host "==================================================" -ForegroundColor Cyan
Write-Host "   Windows SMB + Print Tunnel Client Setup" -ForegroundColor Cyan
Write-Host "==================================================" -ForegroundColor Cyan

# Check and install OpenSSH Client
if (-not (Test-OpenSSHClient)) {
    if (-not (Install-OpenSSHClient)) {
        exit 1
    }
} else {
    Write-Status "OpenSSH Client is already installed"
}

# Generate SSH key
New-SSHKey

# Test connection
if (Test-SSHConnection) {
    # Create tunnel starter
    Start-Tunnel
    
    # Map network drive
    Map-NetworkDrive
    
    # Add printer
    Add-NetworkPrinter
    
    # Create management script
    New-ManagementScript
    
    Write-Host "==================================================" -ForegroundColor Green
    Write-Host "           SETUP COMPLETE!" -ForegroundColor Green
    Write-Host "==================================================" -ForegroundColor Green
    Write-Host "To start the tunnel:" -ForegroundColor Yellow
    Write-Host "  Double-click: $env:USERPROFILE\start-tunnel.bat" -ForegroundColor White
    Write-Host "  Or run: .\manage-tunnel.ps1 start" -ForegroundColor White
    Write-Host ""
    Write-Host "Access Files:" -ForegroundColor Yellow
    Write-Host "  \\localhost\$LocalSmbPort" -ForegroundColor White
    Write-Host "  or open Z: drive" -ForegroundColor White
    Write-Host ""
    Write-Host "Printer:" -ForegroundColor Yellow
    Write-Host "  Use 'Remote-Office-Printer' for printing" -ForegroundColor White
    Write-Host ""
    Write-Host "Management:" -ForegroundColor Yellow
    Write-Host "  .\manage-tunnel.ps1 status" -ForegroundColor White
    Write-Host "==================================================" -ForegroundColor Green
} else {
    Write-Error "Setup cannot continue due to connection issues"
}