# Map-SfldShare.ps1 - keep the SFLD share drive (S:) alive in the interactive
# user's session. Run by the 'GE Map SFLD Share' scheduled task as the
# logged-on user (BUILTIN\Users, Limited): at logon, every few minutes, on
# unlock / console connect, and when a network profile connects.
#
# Rules (the drive is never touched while it works):
#   healthy                 -> exit, no changes.
#   no live mapping         -> map now (nothing to interrupt).
#   mapped but unreadable   -> repair only after $ConfirmFailures consecutive
#                              failed probes, and only if TCP 445 answers.
#   mapped to wrong target  -> same confirm-then-repair path.
#
# Creds: HKLM:\SOFTWARE\GE\SFLD\Credentials\* entries that carry DriveLetter +
# ShareName (written by the vendor SetupCredentials package). The password is
# never put on a command line: New-SmbMapping runs in-process and the
# Credential Manager seed uses CredWrite via P/Invoke. The seed is what lets
# Windows' own logon-time restore and the redirector's transparent reconnect
# authenticate, so both stop producing red-X drives.

param(
    [int]$ConfirmFailures = 2,
    [int]$ProbeTimeoutSec = 8,
    [switch]$Now            # repair on first failed probe (manual / logon use)
)

$ErrorActionPreference = 'Continue'

# --- single instance per session ---------------------------------------
$mutex = New-Object System.Threading.Mutex($false, 'Local\GE-MapSfldShare')
if (-not $mutex.WaitOne(0)) { exit 0 }

# --- logging -----------------------------------------------------------
$logDir = 'C:\Logs\SFLD'
try {
    if (-not (Test-Path $logDir)) { New-Item -Path $logDir -ItemType Directory -Force -EA Stop | Out-Null }
    $probeFile = Join-Path $logDir ('.w-' + [guid]::NewGuid().ToString('N'))
    [IO.File]::WriteAllText($probeFile, 'x'); Remove-Item $probeFile -Force -EA SilentlyContinue
} catch {
    $logDir = Join-Path $env:LOCALAPPDATA 'GE\SfldShare'
    if (-not (Test-Path $logDir)) { New-Item -Path $logDir -ItemType Directory -Force -EA SilentlyContinue | Out-Null }
}
$logFile = Join-Path $logDir 'map-share.log'

function Write-MapLog {
    param([string]$Message, [string]$Level = 'INFO')
    $line = '[{0}] [{1}] [{2}] {3}' -f (Get-Date -Format 'yyyy-MM-dd HH:mm:ss'), $env:USERNAME, $Level, $Message
    Add-Content -Path $logFile -Value $line -ErrorAction SilentlyContinue
    Write-Host $line
}

# --- per-user state (fail counter, last state) --------------------------
$stateRoot = 'HKCU:\SOFTWARE\GE\SfldShare'
function Get-State {
    param([string]$Letter)
    $k = Join-Path $stateRoot $Letter
    if (-not (Test-Path $k)) { New-Item -Path $k -Force | Out-Null }
    $p = Get-ItemProperty -Path $k -EA SilentlyContinue
    return @{
        Key       = $k
        FailCount = [int]($p.FailCount)
        LastState = [string]$p.LastState
        CredHash  = [string]$p.CredHash     # SHA256(user + password) last seeded
    }
}
function Set-State {
    param([string]$Key, [int]$FailCount, [string]$LastState, [string]$CredHash = '')
    Set-ItemProperty -Path $Key -Name FailCount -Value $FailCount -Type DWord -EA SilentlyContinue
    Set-ItemProperty -Path $Key -Name LastState -Value $LastState -Type String -EA SilentlyContinue
    Set-ItemProperty -Path $Key -Name LastRun   -Value (Get-Date -Format 's') -Type String -EA SilentlyContinue
    # only overwrite the seeded-credential hash when the caller supplies one
    if ($CredHash) { Set-ItemProperty -Path $Key -Name CredHash -Value $CredHash -Type String -EA SilentlyContinue }
}
function Get-CredHash {
    param([string]$User, [string]$Password)
    $sha = [System.Security.Cryptography.SHA256]::Create()
    $b = $sha.ComputeHash([Text.Encoding]::UTF8.GetBytes("$User`n$Password"))
    $sha.Dispose()
    return ([BitConverter]::ToString($b) -replace '-', '')
}

# --- probes -------------------------------------------------------------
function Test-Tcp445 {
    param([string]$HostName, [int]$TimeoutMs = 3000)
    try {
        $c = New-Object System.Net.Sockets.TcpClient
        $ar = $c.BeginConnect($HostName, 445, $null, $null)
        $ok = $ar.AsyncWaitHandle.WaitOne($TimeoutMs)
        if ($ok) { $c.EndConnect($ar) }
        $c.Close()
        return $ok
    } catch { return $false }
}

# Real read of the drive root in a runspace with a hard timeout. Test-Path on
# a dead redirector connection can block for 30 s or more.
function Test-DriveReadable {
    param([string]$Drive, [int]$TimeoutSec)
    $ps = [powershell]::Create()
    $null = $ps.AddScript({
        param($d)
        try {
            $e = [System.IO.Directory]::EnumerateFileSystemEntries("$d\").GetEnumerator()
            $null = $e.MoveNext()
            'OK'
        } catch { 'ERR: ' + $_.Exception.GetType().Name + ': ' + $_.Exception.Message }
    }).AddArgument($Drive)
    $h = $ps.BeginInvoke()
    if ($h.AsyncWaitHandle.WaitOne($TimeoutSec * 1000)) {
        $r = $ps.EndInvoke($h); $ps.Dispose()
        return [string]$r[0]
    }
    try { $ps.Stop() } catch {}
    $ps.Dispose()
    return 'TIMEOUT'
}

# --- credential manager seed (CredWrite, no command line) ---------------
$script:credTypeLoaded = $false
function Set-StoredDomainCredential {
    param([string]$Target, [string]$User, [string]$Password)
    if (-not $script:credTypeLoaded) {
        Add-Type -TypeDefinition @'
using System;
using System.Runtime.InteropServices;
public static class GeCredStore {
    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct CREDENTIAL {
        public uint Flags; public uint Type; public string TargetName; public string Comment;
        public System.Runtime.InteropServices.ComTypes.FILETIME LastWritten;
        public uint CredentialBlobSize; public IntPtr CredentialBlob; public uint Persist;
        public uint AttributeCount; public IntPtr Attributes; public string TargetAlias; public string UserName;
    }
    [DllImport("advapi32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
    static extern bool CredWrite([In] ref CREDENTIAL cred, uint flags);
    public static int Write(string target, string user, string password) {
        byte[] blob = System.Text.Encoding.Unicode.GetBytes(password);
        CREDENTIAL c = new CREDENTIAL();
        c.Type = 2;            // CRED_TYPE_DOMAIN_PASSWORD
        c.Persist = 2;         // CRED_PERSIST_LOCAL_MACHINE
        c.TargetName = target; c.UserName = user; c.Comment = "GE SFLD share";
        c.CredentialBlobSize = (uint)blob.Length;
        c.CredentialBlob = Marshal.AllocHGlobal(blob.Length);
        Marshal.Copy(blob, 0, c.CredentialBlob, blob.Length);
        bool ok = CredWrite(ref c, 0);
        int err = ok ? 0 : Marshal.GetLastWin32Error();
        Marshal.FreeHGlobal(c.CredentialBlob);
        return err;
    }
}
'@ -ErrorAction Stop
        $script:credTypeLoaded = $true
    }
    return [GeCredStore]::Write($Target, $User, $Password)
}

# --- mapping ------------------------------------------------------------
function New-Mapping {
    param([string]$Drive, [string]$Unc, [string]$User, [string]$Password, [string]$HostName, [bool]$PreFlight = $false)
    $letter = $Drive.TrimEnd(':')

    # When a mapping (wrong or unreadable) is being replaced, prove the new
    # one can be made BEFORE deleting the old: map the target to a spare
    # letter, then drop it. A failed pre-flight leaves the old mapping in
    # place. 1219 is the exception - it means the old mapping's own session
    # is what blocks the new identity, so deleting it is the only way.
    if ($PreFlight) {
        $used = @((Get-PSDrive -PSProvider FileSystem -EA SilentlyContinue).Name) + @((Get-SmbMapping -EA SilentlyContinue).LocalPath | ForEach-Object { $_.TrimEnd(':') })
        $spare = $null
        foreach ($c in [char[]]'ZYXVUTQPONMLKJIHG') { if ($used -notcontains [string]$c) { $spare = "${c}:"; break } }
        if ($spare) {
            try {
                New-SmbMapping -LocalPath $spare -RemotePath $Unc -UserName $User -Password $Password -Persistent $false -ErrorAction Stop | Out-Null
                $ok = Test-DriveReadable -Drive $spare -TimeoutSec 8
                & net use $spare /delete /y 2>&1 | Out-Null
                if ($ok -ne 'OK') { return "pre-flight via $spare mapped but unreadable ($ok) - old $Drive left in place" }
            } catch {
                $m = $_.Exception.Message
                if ($m -notmatch '1219|more than one user name|Multiple connections') {
                    return "pre-flight via $spare failed: $m - old $Drive left in place"
                }
                Write-MapLog "pre-flight hit 1219 - the existing $Drive session blocks the new identity; replacing it" 'WARN'
            }
        }
    }

    & net use $Drive /delete /y 2>&1 | Out-Null
    Remove-Item -Path "HKCU:\Network\$letter" -Recurse -Force -EA SilentlyContinue
    try {
        New-SmbMapping -LocalPath $Drive -RemotePath $Unc -UserName $User -Password $Password -Persistent $true -ErrorAction Stop | Out-Null
        return $null
    } catch {
        $msg = $_.Exception.Message
        if ($msg -match '1219|more than one user name|Multiple connections') {
            # stale session to the same host under another identity pins the
            # credential; the hidden IPC$ connection is the pin.
            Write-MapLog "1219 credential conflict on $HostName - dropping IPC`$ and retrying" 'WARN'
            & net use "\\$HostName\IPC`$" /delete /y 2>&1 | Out-Null
            try {
                New-SmbMapping -LocalPath $Drive -RemotePath $Unc -UserName $User -Password $Password -Persistent $true -ErrorAction Stop | Out-Null
                return $null
            } catch { return $_.Exception.Message }
        }
        return $msg
    }
}

# --- main ---------------------------------------------------------------
$credsBase = 'HKLM:\SOFTWARE\GE\SFLD\Credentials'
if (-not (Test-Path $credsBase)) { Write-MapLog 'No HKLM SFLD credentials yet - exiting'; exit 0 }

$entries = @()
foreach ($entry in (Get-ChildItem -Path $credsBase -ErrorAction SilentlyContinue)) {
    $p = Get-ItemProperty -Path $entry.PSPath -ErrorAction SilentlyContinue
    if (-not $p -or -not $p.TargetHost -or -not $p.Username -or -not $p.Password) { continue }
    $drive = [string]$p.DriveLetter; $share = [string]$p.ShareName
    if ([string]::IsNullOrWhiteSpace($drive) -or [string]::IsNullOrWhiteSpace($share)) { continue }
    $entries += [pscustomobject]@{
        Drive = $drive.Trim().TrimEnd(':').ToUpper() + ':'
        Unc   = '\\{0}\{1}' -f $p.TargetHost, $share.TrimStart('\').TrimEnd('\')
        Host  = [string]$p.TargetHost
        User  = [string]$p.Username
        Pass  = [string]$p.Password
    }
}
if ($entries.Count -eq 0) { exit 0 }

$exit = 0
foreach ($e in $entries) {
    $letter = $e.Drive.TrimEnd(':')
    $st = Get-State -Letter $letter

    $mapping = Get-SmbMapping -LocalPath $e.Drive -ErrorAction SilentlyContinue
    $present = [bool]$mapping
    $target  = if ($mapping) { [string]$mapping.RemotePath } else { '' }
    $wrong   = $present -and ($target.TrimEnd('\') -ne $e.Unc)

    # Seed Credential Manager whenever the HKLM credential differs from what
    # this profile last seeded (first run, or a password rotation), and again
    # after every repair. Keeps Windows' own logon-time restore authenticating.
    $credHash = Get-CredHash -User $e.User -Password $e.Pass
    if ($st.CredHash -ne $credHash) {
        $rc = Set-StoredDomainCredential -Target $e.Host -User $e.User -Password $e.Pass
        if ($rc -eq 0) {
            Write-MapLog "Seeded Credential Manager for $($e.Host) as $($e.User)$(if ($st.CredHash) { ' (credential changed)' })"
            Set-State -Key $st.Key -FailCount $st.FailCount -LastState $st.LastState -CredHash $credHash
        } else { Write-MapLog "CredWrite for $($e.Host) failed (win32 $rc)" 'WARN' }
    }

    $probe = 'SKIPPED'
    if ($present -and -not $wrong) { $probe = Test-DriveReadable -Drive $e.Drive -TimeoutSec $ProbeTimeoutSec }

    if ($present -and -not $wrong -and $probe -eq 'OK') {
        if ($st.LastState -ne 'healthy') { Write-MapLog "$($e.Drive) healthy -> $($e.Unc)" }
        Set-State -Key $st.Key -FailCount 0 -LastState 'healthy'
        continue
    }

    # --- not healthy --------------------------------------------------
    if (-not $present) {
        Write-MapLog "$($e.Drive) has no live mapping - mapping now"
    } elseif ($wrong) {
        $st.FailCount++
        Write-MapLog "$($e.Drive) points at '$target' not '$($e.Unc)' ($($st.FailCount)/$ConfirmFailures)" 'WARN'
    } else {
        $st.FailCount++
        Write-MapLog "$($e.Drive) mapped (status $($mapping.Status)) but probe failed: $probe ($($st.FailCount)/$ConfirmFailures)" 'WARN'
    }

    if ($present -and -not $Now -and $st.FailCount -lt $ConfirmFailures) {
        Set-State -Key $st.Key -FailCount $st.FailCount -LastState 'suspect'
        continue
    }

    if (-not (Test-Tcp445 -HostName $e.Host)) {
        Write-MapLog "$($e.Host):445 unreachable - not repairing $($e.Drive) this run" 'WARN'
        Set-State -Key $st.Key -FailCount $st.FailCount -LastState 'unreachable'
        $exit = 2
        continue
    }

    Write-MapLog "Repairing $($e.Drive) -> $($e.Unc)$(if ($present) { ' (pre-flight first, existing mapping kept on failure)' })"
    $err = New-Mapping -Drive $e.Drive -Unc $e.Unc -User $e.User -Password $e.Pass -HostName $e.Host -PreFlight $present
    if ($err) {
        Write-MapLog "Repair FAILED $($e.Drive): $err" 'ERROR'
        Set-State -Key $st.Key -FailCount $st.FailCount -LastState 'failed'
        $exit = 1
        continue
    }
    $verify = Test-DriveReadable -Drive $e.Drive -TimeoutSec $ProbeTimeoutSec
    if ($verify -eq 'OK') {
        Write-MapLog "Repaired $($e.Drive) -> $($e.Unc) (verified readable)"
        # re-seed after every repair so the stored credential matches what
        # just proved to work
        $rc = Set-StoredDomainCredential -Target $e.Host -User $e.User -Password $e.Pass
        if ($rc -ne 0) { Write-MapLog "CredWrite after repair failed (win32 $rc)" 'WARN' }
        Set-State -Key $st.Key -FailCount 0 -LastState 'healthy' -CredHash $credHash
    } else {
        Write-MapLog "Mapped $($e.Drive) but verify failed: $verify" 'ERROR'
        Set-State -Key $st.Key -FailCount $st.FailCount -LastState 'failed'
        $exit = 1
    }
}

$mutex.ReleaseMutex() | Out-Null
exit $exit