#Requires -RunAsAdministrator
[CmdletBinding()]
param(
  [switch]$Stop,
  [string]$Socks = "127.0.0.1:59999",
  [string[]]$Direct = @("117.143.9.6"),
  [string]$Dns = "127.0.0.1",
  [switch]$EnableDoH    # 可选：为 $Dns 启用 DoH（Win11+）
)

# --- 全局设置 ---
$ErrorActionPreference = "Stop"
$TunIP   = "192.168.123.1"
$TunMask = "255.255.255.0"
$State   = Join-Path $PSScriptRoot ".wintun-state.json"

# --- 常量：防火墙规则名 ---
$FwNames = @(
  "Block NonTUN DNS TCPv4",
  "Block NonTUN DNS UDPv4",
  "Block NonTUN DNS TCPv6",
  "Block NonTUN DNS UDPv6"
)

function Get-Uplink {
  <#
    选一个物理上行（默认路由优先，排除虚拟网卡）
    Returns: [pscustomobject] @{ Alias; Index; Gateway }
  #>
  $bad = "wintun|wireguard|tap|vEthernet|Hyper-V|VirtualBox|VMware|Npcap|Loopback|Docker|Tailscale|ZeroTier|Hamachi"
  $routes = Get-NetRoute -DestinationPrefix "0.0.0.0/0" -AddressFamily IPv4 |
            Sort-Object RouteMetric, InterfaceMetric
  foreach ($r in $routes) {
    $if = Get-NetIPInterface -InterfaceIndex $r.InterfaceIndex -ErrorAction SilentlyContinue |
            Where-Object { $_.AddressFamily -eq 2 } # 2 means IPv4
    if ($null -ne $if -and ($if.InterfaceAlias -notmatch $bad)) {
      return [pscustomobject]@{ Alias = $if.InterfaceAlias; Index = $if.InterfaceIndex; Gateway = $r.NextHop }
    }
  }
  throw "No physical uplink found."
}

function Remove-DnsEnforcement {
  <#
    清理：NRPT + 防火墙规则
  #>
  try {
    Get-DnsClientNrptRule -ErrorAction SilentlyContinue |
      Where-Object { $_.Namespace -eq "." } |
      Remove-DnsClientNrptRule -ErrorAction SilentlyContinue -Force | Out-Null
  } catch {}

  foreach ($n in $FwNames) {
    Get-NetFirewallRule -DisplayName $n -ErrorAction SilentlyContinue | Remove-NetFirewallRule -ErrorAction SilentlyContinue
  }
}

function Apply-DnsEnforcement {
  param(
    [Parameter(Mandatory=$true)][string[]]$DnsServers,
    [Parameter(Mandatory=$true)][int]$TunIfIndex
  )
  <#
    强制 DNS：
    1) NRPT "." 指向 $DnsServers
    2) (已禁用) 封非 TUN 的 53 (TCP/UDP, v4/v6)
  #>
  # 清理旧 NRPT
  Remove-DnsEnforcement

  # NRPT：所有域名都走指定 DNS
  Add-DnsClientNrptRule -Namespace "." -NameServers $DnsServers -Comment "Force all DNS via TUN" | Out-Null

  # --- 防火墙规则部分已移除，以解决兼容性问题 ---
  
}

function Enable-DoHIfRequested {
  param(
    [Parameter(Mandatory=$true)][string]$Server
  )
  <#
    为指定 IPv4 DNS 启用 DoH（Win11+）。失败时静默跳过。
  #>
  if (-not $EnableDoH) { return }
  if ($Server -match '^\d{1,3}(\.\d{1-3}){3}$') {
    try {
      # Cloudflare 模板示例；如换 DNS，请替换模板
      netsh dns add encryption server=$Server dohtemplate=https://cloudflare-dns.com/dns-query autoupgrade=yes udpfallback=no | Out-Null
    } catch {}
  }
}

function Start-Tun {
  <#
    启动 tun2socks，创建 TUN，设置地址与 DNS 强制策略，添加默认路由与直连例外
  #>
  $t2s = Join-Path $PSScriptRoot "tun2socks.exe"
  $dll = Join-Path $PSScriptRoot "wintun.dll"
  if (-not (Test-Path $t2s)) { throw "Missing tun2socks.exe" }
  if (-not (Test-Path $dll)) { throw "Missing wintun.dll (same folder as script)" }

  # 先清掉历史 tun2socks
  Get-Process tun2socks -ErrorAction SilentlyContinue | Stop-Process -Force -ErrorAction SilentlyContinue

  $uplink = Get-Uplink
  Write-Host ("Uplink: {0}  Gateway: {1}" -f $uplink.Alias, $uplink.Gateway)

  $before = (Get-NetIPInterface | Where-Object { $_.AddressFamily -eq 2 }).InterfaceIndex

  # 启动 tun2socks
  Start-Process -FilePath $t2s -ArgumentList @(
    "-device","wintun",
    "-proxy",("socks5://{0}" -f $Socks),
    "-interface",$uplink.Alias
  ) -WindowStyle Hidden | Out-Null

  # 识别新出现的 TUN 接口
  $tun = $null
  for ($i=0; $i -lt 60; $i++) {
    Start-Sleep -Milliseconds 200
    $after = Get-NetIPInterface | Where-Object { $_.AddressFamily -eq 2 }
    $new   = $after | Where-Object { $before -notcontains $_.InterfaceIndex }
    if ($new) { $tun = $new | Sort-Object InterfaceMetric | Select-Object -First 1; break }
  }
  if (-not $tun) {
    $tun = Get-NetIPInterface | Where-Object { $_.AddressFamily -eq 2 } |
           Where-Object { $_.InterfaceAlias -match "wintun|tun|wireguard" } |
           Sort-Object InterfaceMetric | Select-Object -First 1
  }
  if (-not $tun) { throw "Wintun interface not found." }

  # 配置 TUN 地址
  netsh interface ipv4 set address name="$($tun.InterfaceAlias)" source=static addr=$TunIP mask=$TunMask

  # 设 TUN 的 DNS (使用 netsh 以增强兼容性)
  $dnsServers = @($Dns)
  netsh interface ipv4 set dnsservers name="$($tun.InterfaceAlias)" static $($dnsServers[0]) primary

  # 可选启用 DoH
  Enable-DoHIfRequested -Server $Dns

  # 强制：NRPT + 防火墙
  Apply-DnsEnforcement -DnsServers $dnsServers -TunIfIndex $tun.InterfaceIndex

  # 把默认路由指向 TUN（所有流量含 DNS 都从隧道走）
  route -p add 0.0.0.0 mask 0.0.0.0 $TunIP metric 1 if $($tun.InterfaceIndex) | Out-Null

  # 直连例外（这些目标直接走物理上行）
  foreach ($ip in $Direct) {
    route -p add $ip mask 255.255.255.255 $($uplink.Gateway) if $($uplink.Index) metric 5 | Out-Null
  }

  # 降低 TUN 接口度量 (已注释掉，以解决兼容性问题)
  # try { Set-NetIPInterface -InterfaceIndex $tun.InterfaceIndex -InterfaceMetric 5 -ErrorAction SilentlyContinue } catch {}

  # 记录状态
  @{
    TunAlias = $tun.InterfaceAlias
    TunIndex = $tun.InterfaceIndex
    TunIP    = $TunIP
    Direct   = $Direct
    Dns      = $dnsServers
  } | ConvertTo-Json | Set-Content -Path $State -Encoding UTF8

  Write-Host ("TUN ready: {0} -> SOCKS {1}; DNS {2}; direct: {3}" -f $tun.InterfaceAlias, $Socks, ($dnsServers -join ","), ($Direct -join ", ")) -ForegroundColor Green
}

function Stop-Tun {
  <#
    清理默认路由、直连路由、NRPT、防火墙、进程与状态
  #>
  $tunIPLocal   = $TunIP
  $directLocal  = $Direct
  if (Test-Path $State) {
    try {
      $s = Get-Content $State | ConvertFrom-Json
      if ($s.TunIP)   { $tunIPLocal  = $s.TunIP }
      if ($s.Direct)  { $directLocal = $s.Direct }
    } catch {}
  }

  # 路由清理
  route delete 0.0.0.0 mask 0.0.0.0 $tunIPLocal | Out-Null
  foreach ($ip in $directLocal) { route delete $ip | Out-Null }

  # 清理 DNS 强制策略
  Remove-DnsEnforcement

  # 结束进程与状态文件
  Get-Process tun2socks -ErrorAction SilentlyContinue | Stop-Process -Force -ErrorAction SilentlyContinue
  Remove-Item $State -ErrorAction SilentlyContinue | Out-Null

  Write-Host "Stopped and cleaned routes & DNS policies." -ForegroundColor Yellow
}

# --- 入口 ---
try {
  if ($Stop) { Stop-Tun } else { Start-Tun }
} catch {
  Write-Host ("Error: {0}" -f $_.Exception.Message) -ForegroundColor Red
  exit 1
}