#Requires -Modules powershell-yaml
param (
    [string]$Chain = $env:CHAIN,
    [string]$Stacks = $env:STACKS
)
$ErrorActionPreference = 'Stop'

function Start-IptablesProcess {
    param (
        [array]$ArgumentList
    )

    $splat = @{
        FilePath     = 'iptables'
        ArgumentList = $ArgumentList
        Wait         = $true
        PassThru     = $true
    }

    Start-Process @splat
}
function Test-IptablesChain {
    param (
        [string]$Chain,
        [string]$Table = 'nat'
    )

    $output = iptables -t $Table -S $Chain
    $reference = '-N {0}' -f $Chain

    #check $output in case $Chain has no rules
    #otherwise check the first line with $output[0]
    $output.Count -gt 0 `
    -and ($output -eq $reference `
          -or `
          $output[0] -eq $reference)
}
function Add-IptablesChain {
    param (
        [string]$Chain,
        [string]$Table = 'nat'
    )

    if(-not (Test-IptablesChain -Chain $Chain)) {
        iptables -t $Table -N $Chain
    }
}
function Test-IptablesRule {
    param (
        [string]$Chain,
        [string]$Table = 'nat',
        [array]$Rule
    )

    $argument_list = @(
        '-t'
        $Table
        '-C'
        $Chain
    )+$Rule

    $check = Start-IptablesProcess -ArgumentList $argument_list

    Write-Output ($check.ExitCode -eq 0)
}
function Add-IptablesRule {
    param (
        [string]$Chain,
        [string]$Table = 'nat',
        [array]$Rule
    )

    $argument_list = @(
        '-t'
        $Table
        '-I'
        $Chain
    )+$Rule

    if(-not (Test-IptablesRule -Chain $Chain -Rule $Rule)) {
        $add = Start-IptablesProcess -ArgumentList $argument_list
        if(0 -eq $add.ExitCode) {
            Write-Output $true
        } else {
            Write-Error 'Adding iptables rule failed'
        }
    }
}
function Get-DockerIngressAddress {
    param (
        [string]$BridgeDevice = 'docker_gwbridge'
    )

    $bridge = docker network inspect $BridgeDevice | ConvertFrom-Json
    $bridge.Containers.{ingress-sbox}.IPv4Address -replace '/.*'
}

# setup SWARM-NAT chain
Write-Output ('Create chain {0}' -f $Chain)
Add-IptablesChain -Chain $Chain
$chain_rule = @(
    '-m','addrtype'
    '--dst-type','LOCAL'
    '-j',$Chain
)
Add-IptablesRule -Chain 'PREROUTING' -Rule $chain_rule
Add-IptablesRule -Chain 'OUTPUT' -Rule $chain_rule


$ingress_address = Get-DockerIngressAddress


foreach($yaml in (Get-ChildItem -Path $Stacks -Include 'docker-compose.yml' -Recurse)) {
    Write-Output ('Processing {0}' -f $yaml)
    $definition = Get-Content -Path $yaml -Raw | ConvertFrom-Yaml

    foreach($port in $definition.services.Values.ports) {
        $nat = @{
            protocol = $null
            ip = $null
            port = $null
            published_port = $null
        }

        if($port.Count -eq 4) {
            #long form
            $published_splitted = $port.published -split ':'

            $nat.protocol = $port.protocol
            $nat.ip = $published_splitted[0]
            $nat.port = $port.target
            $nat.published_port = $published_splitted[1]
        } else {
            #short form
            $all_splitted = $port -split ':'
            $port_splitted = $all_splitted[2] -split '/'

            $nat.protocol = $port_splitted[1]
            $nat.ip = $all_splitted[0]
            $nat.port = $port_splitted[0]
            $nat.published_port = $all_splitted[1]
        }

        if(!$nat.protocol) {
            #this is also Docker's default
            $nat.protocol = 'tcp'
        }
        if($nat.Values -contains $null) {
            #if $nat doesn't conatain all needed attributes skip the port
            $error_message = 'Skipping port, because $nat contains $null: {0}' `
                             -f ($nat | ConvertTo-Json)
            Write-Error -Message $error_message `
                        -ErrorAction Continue
            continue
        }

        if($nat.published_port -ne $nat.port) {
            Write-Output ('Additional NAT rule required, because published_port {0} and target {1} differ' `
                          -f $nat.published_port,$nat.port)
            Write-Output ('Add rule for {0}:{1}' -f $nat.ip,$nat.port)
            
            $rule = @(
                '-p',$nat.protocol
                '-m',$nat.protocol
                '--destination',$nat.ip
                '--dport',$nat.port
                '-j','DNAT'
                '--to-destination','"{0}:{1}"' -f $ingress_address,$nat.published_port
            )
            Add-IptablesRule -Chain $Chain -Rule $rule
        } else {
            Write-Output ('No additional rule needed for {0}:{1}' -f $nat.ip,$nat.port)
        }
    }
}