# See: https://aws.amazon.com/about-aws/whats-new/2022/05/aws-systems-manager-support-port-forwarding-remote-hosts-using-session-manager/
# And: https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-sessions-start.html#sessions-remote-port-forwarding
function ssm_tunnel --description 'Access private AWS ressources via native SSM port forwarding through bastion'

    if test (count $argv) -lt 2
        echo "Provide an Host:Port Mapping & an Local Port"
        return
    end

    if [ "$env" != "ops" ]
        echo "Switch to OPS ENV first"
        return
    end

    # Get bastion host ID using it's name
    set INSTANCE_ID (aws ec2 describe-instances \
            --filter 'Name=tag:Name,Values=INSTANCE_NAME_HERE'\
            --query "Reservations[].Instances[?State.Name == 'running'].InstanceId[]"\
            --output text)
   
    # AWS Remote host i.e internal.s.twaice / dualstack.internal-stg-s-inter-9ov8h1o4saoa-793915940.eu-west-1.elb.amazonaws.com
    set REMOTE_HOST (string split -f1 : $argv[1])
    set REMOTE_PORT (string split -f2 : $argv[1])
    # Local port to bind for forwarding
    set LOCAL_PORT $argv[2]

     # Seconds until forwarding session times out (6h max)
    set TIMEOUT 21600

    # Start tunnel session
    echo "Starting Tunnel"
    aws ssm start-session --target $INSTANCE_ID \
    --document-name "AWS-StartPortForwardingSessionToRemoteHost" \
    --parameters "host=[$REMOTE_HOST],portNumber=[$REMOTE_PORT],localPortNumber=[$LOCAL_PORT]"

    echo "Tunneling Session Exited."
end

# This function is for Platform Admin / Infrastructure users only.
# @TODO: It still utilized the old style of ssm tunneling -> needs c&p of new style as well
function ssm_tunnel_admin --description 'access private AWS ressources via Bastion Host'

    if test (count $argv) -lt 2
        echo "Provide an Host:Port Mapping & an Local Port"
        return
    end

    if [ "$env" != "ops" ]
        echo "Switch to OPS ENV first"
        return
    end

    # Get bastion host ID using it's name
    set INSTANCE_ID (aws ec2 describe-instances \
            --filter 'Name=tag:Name,Values=INSTANCE_NAME_HERE'\
            --query "Reservations[].Instances[?State.Name == 'running'].InstanceId[]"\
            --output text)
    
    # AWS Remote host i.e elastic.aws.com:80
    set REMOTE_HOST $argv[1]

    # Ports to bind for forwarding
    set LOCAL_PORT $argv[2]
    set REM_PORT (math (random) % 65535 + 2000)
    set TIMEOUT 21600 #seconds until forwarding session times out (48h max)

    # make sure jq is installed
    # Start socat on the remote server
    set CMD "'sudo socat TCP4-LISTEN:$REM_PORT,reuseaddr,fork TCP4:$REMOTE_HOST'"
    set CMD_INVOC_ID (aws ssm send-command --instance-ids $INSTANCE_ID \
    --document-name 'AWS-RunShellScript' \
    --parameters "commands=$CMD,executionTimeout=$TIMEOUT" --output json | jq -r '.Command.CommandId')
    
    # Start tunnel session
    echo "Starting Tunnel"
    aws ssm start-session --target $INSTANCE_ID \
    --document-name "AWS-StartPortForwardingSession" \
    --parameters "portNumber=$REM_PORT,localPortNumber=$LOCAL_PORT"

    # stop socat on the bastion
    aws ssm cancel-command --command-id $CMD_INVOC_ID
    echo "Command Cancelled Successfully."
end