# auto-pf.ps1
#
# Windowsでsshのポートフォワーディングを維持するための、
# PortForwarderのような使い勝手とautosshのような再接続機能を持つ、
# Windowsの標準機能だけで実装したPowerShellスクリプト。
# Copyright (C) 2023 flatray.com
#
# Base Code:
# https://qiita.com/magiclib/items/cc2de9169c781642e52d
#
#------------------------------------------------------------------------------
# This software is released under the MIT License.
# ライセンス全文: https://opensource.org/licenses/MIT
# 日本語訳: https://licenses.opensource.jp/MIT/MIT.html

param(
    ## 接続先をコンフィグファイルから:
    #   $True = 読み込む(デフォルト)
    #   $False = 読み込まない(sshの設定を含めてプログラム内で指定する)
    [Bool]$auto = $True,
    ## 接続先を記載したコンフィグファイル(≒ssh設定ファイル)
    #   デフォルトは c:\Users\ユーザ名\.ssh\config
    [String]$config = $env:USERPROFILE + "\.ssh\config",
    ## ツールチップの日時書式 [ (Get-Date).ToString(...) の ... の部分 ]
    #   または notime (日時は表示しない) / none (プログラム名のみ表示する)
    [String]$tooltip = "MM/dd HH:mm",
    ## 接続先をコンフィグファイルから読み出す際のキーワード
    [String]$keyword = "auto-pf",
    ## 接続時にバルーンを表示する (-baloon 1 と指定する)
    [Bool]$baloon = $False,
    ## 再接続(ジョブチェック)の実行間隔(秒)
    [Int]$wait = 5,
    ## デバッグメッセージを出力
    [Bool]$debug = $False,
    ## PIDファイル
    [String]$pidfile = $env:USERPROFILE + "\auto-pf.pid.txt",
    ## ログファイル
    [String]$logfile = $env:USERPROFILE + "\auto-pf.log.txt",
    ## エラーログファイル
    [String]$errorfile = $env:USERPROFILE + "\auto-pf.error.txt",
    ## ホストを引数で指定 ([Name:]Host,[Name:]Host,...)
    #   指定しない場合はコンフィグファイルまたはプログラム内設定を使用
    #   これを使用する時に -auto 0 と設定しないでください
    [String]$connect = "",
    ## sshコマンドとオプション
    [String]$ssh = "ssh.exe -A -N -T"
     )

##-----------------------------------------------------------------------------
## カスタマイズ (sshの設定をプログラム内に書く場合) start ↓↓↓↓↓↓↓↓
# デフォルト値
$DEFAULT_USER = "YourDefaultUserName"
$DEFAULT_COMMAND = "ssh.exe"
$DEFAULT_PORT = "22"
$DEFAULT_FWD = ""

# 以下のコメントを外し必要な数だけ接続先を記述してください。
# 1つでよければ減らし、3つ以上必要なら増やしてください。
# 必須: name, dest
# 任意: command, user, port, fwd (省略時は上記の $DEFAULT_* を適用)

# $REMOTE1 = @{
#     name = "Remote1";
#     dest = "RemoteServer1";
# #    command = "ssh.exe";
# #    user = "RemoteUser1";
# #    port = "22";
# #    fwd = "-L 5901:localhost:5901"
# }
# $REMOTE2 = @{
#     name = "Remote2";
#     dest = "RemoteServer2";
# #    command = "ssh.exe";
# #    user = "RemoteUser2";
# #    port = "22";
# #    fwd = "-L 5902:localhost:5901"
# }

# 以下のコメントを外し接続先をリストに登録してください。
# 1つでよければ減らし、3つ以上必要なら増やしてください。

# $REMOTE_LIST = @($REMOTE1, $REMOTE2)

## カスタマイズ end ↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑
##-----------------------------------------------------------------------------

## 変数
# timer_function実行間隔(ミリ秒)
$TIMER_INTERVAL = $wait * 1000
# ジョブ情報
$global:job = @{}
# 接続回数
$global:count = @{}
# 最後に接続した日時(ツールチップ用)
$global:last_connect = @{}
# このスクリプト名
$this_program = $myInvocation.MyCommand.name
# このスクリプトのフルパス
$this_path = $myInvocation.MyCommand.Path
# 多重起動チェック用
$MUTEX_NAME = "Global\" + $this_program

Add-Type -AssemblyName System.Windows.Forms
Add-Type -AssemblyName System.Drawing

# オプション -connect から接続先を得る
if ($connect -ne "") {
    $REMOTE_LIST = @()
    foreach ($pair in $connect -split ",") {
	#Write-Host $pair
	if ($pair.Contains(":")) {
	    $p = $pair -split ":"
	    $name = $p[0]
	    $dest = $p[1]
	}
	else {
	    $name = $pair
	    $dest = $pair
	}
	$REMOTE_LIST += @{ name = $name; dest = $dest }
    }
}
# コンフィグファイルから接続先を読み出す
elseif ($auto) {
    $REMOTE_LIST = @()
    $conf = (Get-Content $config) -as [string[]]
    for ($i = 0; $i -lt $conf.count; ++$i) {
	$line = $conf[$i]
	if ($line -match "#\s*$keyword\s+(?<name>\S+)") {
	    $name = $Matches.name
	    $line = $conf[$i + 1]
	    $null = $line -match "Host\s+(?<dest>\S+)"
	    if ($Matches) {
		$dest = $Matches.dest
		$REMOTE_LIST += @{ name = $name; dest = $dest }
	    }
	}
	elseif ($line -match "#\s*$keyword") {
	    $line = $conf[$i + 1]
	    $null = $line -match "Host\s+(?<dest>\S+)"
	    if ($Matches) {
		$dest = $Matches.dest
		$REMOTE_LIST += @{ name = $dest; dest = $dest }
	    }
	}
    }
    if (-not $REMOTE_LIST) {
	Write-Host "コンフィグファイル $config から接続先一覧を取得できません"
	Read-Host "エンターキーを押すと終了します"
	exit 1
    }
}

# sshジョブを起動する
function start_job ($remote) {
    # 自動モード (設定はすべて .ssh/config に記載)
    if ($auto) {
	$dest = $remote.dest
	$sb = [scriptblock]::Create("$ssh $dest")
    }
    # 手動モード (設定はプログラム内に記載)
    else {
	$dest = $remote.dest
	if ($remote.ContainsKey("command")) {
	    $command = $remote.command
	}
	else {
	    $command = $DEFAULT_COMMAND
	}
	if ($remote.ContainsKey("user")) {
	    $user = $remote.user
	}
	else {
	    $user = $DEFAULT_USER
	}
	if ($remote.ContainsKey("port")) {
	    $port = $remote.port
	}
	else {
	    $port = $DEFAULT_PORT
	}
	if ($remote.ContainsKey("fwd")) {
	    $fwd = $remote.fwd
	}
	else {
	    $fwd = $DEFAULT_FWD
	}
	$sb = [scriptblock]::Create("$command -p $port -l $user -T -A -N $fwd $dest")
    }
    return Start-Job -ScriptBlock $sb
}

# ツールチップのメッセージを設定する
function set_tooltip ($notify) {
    $text = @()
    if ($tooltip -eq "none") {
	$text += $this_program
    }
    else {
	foreach ($remote in $REMOTE_LIST) {
	    $name = $remote.name
	    if ($global:last_connect.ContainsKey($name)) {
		if ($global:last_connect.$name -eq "") {
		    $t = $name + " " + $global:count.$name
		}
		else {
		    $t = $global:last_connect.$name + " " + $name + " " + $global:count.$name
		}
		$text += $t
	    }
	}
    }
    $notify.Text = $text -join "`n"
}

# ログファイルのヘッダ出力
function log_header ($tag) {
    ("===== " + $tag + " =====") | Out-File -Append $logfile
    (Get-Date).ToString("yyyy/MM/dd HH:mm:ss") | Out-File -Append $logfile
}
function error_header ($tag) {
    ("===== " + $tag + " =====") | Out-File -Append $errorfile
    (Get-Date).ToString("yyyy/MM/dd HH:mm:ss") | Out-File -Append $errorfile
}

# ここに定期実行する処理を実装
function timer_function($notify){
    foreach ($remote in $REMOTE_LIST) {
	$name = $remote.name

	# 初回起動時
	if (! $global:job.$name) {
	    # 日付
	    $datetime_long = (Get-Date).ToString("yyyy-MM-dd HH:mm:ss")
	    if ($tooltip -eq "none" -or $tooltip -eq "notime") {
		$datetime_short = ""
	    }
	    else {
		$datetime_short = (Get-Date).ToString($tooltip)
	    }
	    # ジョブ起動
	    $global:job.$name = start_job($remote)
	    # 接続回数を初期化
	    $global:count.$name = 1
	    # コンソールにメッセージを出力
	    $message = "$datetime_long $name Start."
	    Write-Host $message
	    # 最終接続日時を更新
	    $global:last_connect.$name = $datetime_short
	    # ツールチップを設定
	    set_tooltip($notify)
	    # バルーンで表示
	    if ($baloon) {
		$notify.BalloonTipIcon = 'Info'
		$notify.BalloonTipText = $message
		$notify.BalloonTipTitle = $this_program
		$notify.ShowBalloonTip(500)
	    }
	}
	# ジョブを再起動 (状態が Running ではない時、大抵は Completed)
	# (Running, Completed, Failed, etc?)
	elseif ($global:job.$name.State -ne "Running") {
	    $job = $global:job.$name
	    # ステータスを保存
	    $state = $job.State
	    # Completed(ジョブ終了)でない時はジョブを終了する
	    if ($state -ne "Completed") {
		Stop-Job -Job $job
		Wait-Job -Job $job
	    }
	    # 停止理由をログに出力
	    if ($logfile -ne "none") {
		log_header("STOP ssh job")
		Receive-Job -Job $job *>&1 | Format-List -force | Out-File -Append $logfile
	    }
	    # ジョブ一覧からジョブを削除
	    Remove-Job -Job $job

	    # 日付
	    $datetime_long = (Get-Date).ToString("yyyy-MM-dd HH:mm:ss")
	    if ($tooltip -eq "none" -or $tooltip -eq "notime") {
		$datetime_short = ""
	    }
	    else {
		$datetime_short = (Get-Date).ToString($tooltip)
	    }
	    # ジョブ起動
	    $global:job.$name = start_job($remote)
	    # 接続回数を増やす
	    $global:count.$name += 1
	    # コンソールにメッセージを出力
	    $message = "$datetime_long $name Restart by " + $state + ": Count " + $global:count.$name
	    Write-Host $message
	    # 最終接続日時を更新
	    $global:last_connect.$name = $datetime_short
	    # ツールチップを設定
	    set_tooltip($notify)
	    # バルーンで表示
	    if ($baloon) {
		$notify.BalloonTipIcon = 'Info'
		$notify.BalloonTipText = $message
		$notify.BalloonTipTitle = $this_program
		$notify.ShowBalloonTip(500)
	    }
	}
	# ジョブが存在する (デバッグ用)
	else {
	    if ($debug) {
		$datetime_long = (Get-Date).ToString("yyyy-MM-dd HH:mm:ss")
		Write-Host $datetime_long $name "Found job"
	    }
	}
    }
}

## 以下、参考にしたプログラムがベース

# イベントハンドラ内で書き込みたい変数は `script` スコープで宣言する https://github.com/yokra9/RunCat_for_Windows_on_PowerShell/blob/master/RunCatPS/src/runcat.ps1
$script:isTerminalWindowClosed = $True;  # ターミナルウィンドウが閉じているか否か

function main(){
    $mutex = New-Object System.Threading.Mutex($false, $MUTEX_NAME)
    # 多重起動チェック
    $mutex_ok = $True
    if (! $mutex.WaitOne(0, $false)){
	$mutex_ok = $False
	$mutex.Close()
	## 古いプロセスを終了して新規起動する(ハング時など対策)
	# 古いプロセスを探す
	$commandline = ""
	$processid = 0
	# powershell のプロセスを対象に
	foreach ($p in Get-Process powershell) {
	    $id = $p.id
	    # 自分は対象外とする
	    if ($id -eq $pid) {
		continue
	    }
	    # コマンドラインに同じスクリプト名を含むのが対象
	    $q = Get-WmiObject win32_process -filter processid=$id
	    $c = $q.CommandLine
	    if ($c -and $c.ToLower().Contains($this_path.ToLower())) {
		$commandline = $c
		$processid = $id
		break
	    }
	}
	# プロセスがあったら停止する
	if ($commandline -ne "") {
	    Write-Host "このプログラム(ProcessID: $Pid / CommandLine $this_path)と同じプロセスが動作しています。"
	    Write-Host ""
	    Write-Host "ProcessID:" $processid
	    Write-Host "CommandLine:" $commandline
	    Write-Host ""
	    Write-Host "上記のプロセスを終了し新規起動する場合は y を入力してください。"
	    Write-Host "(さもなければ y 以外を入力するか Ctrl+C で終了してください)"
	    Write-Host ""
	    $Input = Read-Host "入力"
	    if ($Input -eq "y") {
		# まず CloseMainWindow で安全に停止を試みる
		$p = Get-Process -Id $processid
		$null = $p.CloseMainWindow()
		Start-Sleep -Second 2
		# それでも止まらなければ Stop-Process を試みる
		$ErrorActionPreference = "SilentlyContinue"
		if (Get-Process -Id $processid) {
		    Stop-Process -Id $processid
		    Start-Sleep -Second 2
		}
		$ErrorActionPreference = "Continue"
	    }
	    else {
		exit
	    }
	}
	# 競合プロセスがないのに何故か mutex をシグナル状態にできない
	else {
	    if ($errorfile -ne "none") {
		error_header("MUTEX1")
		"No process, but WaitOne() failed." | Out-File -Append $errorfile
	    }
	    exit
	}
    }

    # 競合プロセスを停止した時は再度 mutex をシグナル状態にする
    if (! $mutex_ok) {
	$mutex = New-Object System.Threading.Mutex($false, $MUTEX_NAME)
	$mutex_ok = $False
	# 何度か試みる
	foreach ($i in 1..5) {
	    if ($mutex.WaitOne(0, $false)) {
		$mutex_ok = $True
		break
	    }
	    Start-Sleep 1
	}
	# シグナル状態にならない時は終了
	if (! $mutex_ok) {
	    if ($errorfile -ne "none") {
		error_header("MUTEX2")
		("My Pid " + $Pid) | Out-File -Append $errorfile
		foreach ($p in Get-Process powershell) {
		    $id = $p.id
		    ("PID " + $id) | Out-File -Append $errorfile
		    $q = Get-WmiObject win32_process -filter processid=$id
		    ("CommandLine " + $q.CommandLine) | Out-File -Append $errorfile
		}
	    }
	    $mutex.Close()
	    exit
	}
    }

    # PIDを出力
    $pid | Out-File -FilePath $pidfile
    # タスクバー非表示
    $windowcode = '[DllImport("user32.dll")] public static extern bool ShowWindowAsync(IntPtr hWnd, int nCmdShow); [DllImport("user32.dll")] public static extern bool SetForegroundWindow(IntPtr hWnd);'
    $asyncwindow = Add-Type -MemberDefinition $windowcode -name Win32ShowWindowAsync -namespace Win32Functions -PassThru
    $windowHandle = (Get-Process -PID $pid).MainWindowHandle;
    $null = $asyncwindow::ShowWindowAsync($windowHandle, 0)

    $application_context = New-Object System.Windows.Forms.ApplicationContext
    $timer = New-Object Windows.Forms.Timer
    $path = Get-Process -id $pid | Select-Object -ExpandProperty Path # icon用
    # アイコンファイルがあればそれを使う。なければデフォルトを使う。
    $icon_file = $PSScriptRoot + "\auto-pf.ico"
    if (Test-Path -Path $icon_file -PathType Leaf) {
	$icon = New-Object System.Drawing.Icon $icon_file
    }
    else {
	$icon = [System.Drawing.Icon]::ExtractAssociatedIcon($path)
    }
    # タスクトレイアイコン
    $notify_icon = New-Object System.Windows.Forms.NotifyIcon
    $notify_icon.Icon = $icon
    $notify_icon.Visible = $true

    # アイコンクリック時のイベント
    $notify_icon.add_Click({
			       if ($_.Button -eq [Windows.Forms.MouseButtons]::Left) {
				   # ターミナルウィンドウの表示・非表示を切り替える
				   $script:isTerminalWindowClosed = ! $script:isTerminalWindowClosed;
				   $windowMode = if($script:isTerminalWindowClosed) { Write-Output 0; } else { Write-Output 9; };
				   [void]$asyncwindow::ShowWindowAsync($windowHandle, $windowMode);
				   # 最前面表示
				   if (! $script:isTerminalWindowClosed) {
				       [void]$asyncwindow::SetForegroundWindow($windowHandle);
				   }
			       }
			   })

    # メニュー
    $menu_item_exit = New-Object System.Windows.Forms.MenuItem
    $menu_item_exit.Text = $this_program + " Exit"
    $notify_icon.ContextMenu = New-Object System.Windows.Forms.ContextMenu
    $notify_icon.contextMenu.MenuItems.AddRange($menu_item_exit)

    # Exitメニュークリック時のイベント
    $menu_item_exit.add_Click({
				  $application_context.ExitThread()
			      })

    # タイマーイベント.
    $timer.Enabled = $true
    $timer.Add_Tick({
			$timer.Stop()

			timer_function($notify_icon)

			# インターバルを再設定してタイマー再開
			$timer.Interval = $TIMER_INTERVAL
			$timer.Start()
		    })

    $timer.Interval = 1
    $timer.Start()

    [void][System.Windows.Forms.Application]::Run($application_context)

    $timer.Stop()
    $notify_icon.Visible = $false
    $mutex.ReleaseMutex()
    $mutex.Close()
}

try {
    main
}
catch {
    if ($errorfile -ne "none") {
	error_header("main catch")
	$Pid | Out-File -Append $errorfile
	$error[0] | Format-List -force | Out-File -Append $errorfile
    }
}