EonaCat.RdpColorizer/EonaCat.RdpColorizer/RdpSession.cs

using System.Diagnostics;
using System.Text;
using static EonaCat.RdpColorizer.NativeMethods;

namespace EonaCat.RdpColorizer
{
    // This file is part of the EonaCat project(s) which is released under the Apache License.
    // See the LICENSE file or go to https://EonaCat.com/License for full license details.

    /// <summary>
    /// Represents one launched (or attached) Remote Desktop connection whose
    /// title bar is colorized and labeled with an overlay tag. Polls the
    /// target window's position so the tag follows it when moved or resized.
    /// </summary>
    public sealed class RdpSession : IDisposable
    {
        public string Target { get; }
        public string Label { get; }
        public Color TagColor { get; }

        private Process? _process;
        private IntPtr _rdpHwnd = IntPtr.Zero;
        private readonly System.Windows.Forms.Timer _watchTimer;
        private TagOverlayForm? _overlay;
        private bool _disposed;

        public RdpSession(string target, string label, Color tagColor)
        {
            Target = target;
            Label = label;
            TagColor = tagColor;

            _watchTimer = new System.Windows.Forms.Timer { Interval = 150 };
            _watchTimer.Tick += WatchTimer_Tick;
        }

        /// <summary>Launches mstsc.exe against the target and starts tracking it.</summary>
        public void Launch()
        {
            var psi = new ProcessStartInfo
            {
                FileName = "mstsc.exe",
                Arguments = $"/v:{Target}",
                UseShellExecute = true
            };
            _process = Process.Start(psi);
            _watchTimer.Start();
        }

        /// <summary>Attaches to an already-open mstsc window instead of launching a new one.</summary>
        public void AttachToWindow(IntPtr hwnd)
        {
            _rdpHwnd = hwnd;
            OnWindowFound();
            _watchTimer.Start();
        }

        private void WatchTimer_Tick(object? sender, EventArgs e)
        {
            if (_rdpHwnd == IntPtr.Zero)
            {
                if (_process == null)
                {
                    return;
                }

                IntPtr hwnd = FindMstscWindowForProcess(_process.Id);
                if (hwnd != IntPtr.Zero)
                {
                    _rdpHwnd = hwnd;
                    OnWindowFound();
                }
                return;
            }

            if (!IsWindow(_rdpHwnd))
            {
                // The RDP window was closed; clean up and stop watching.
                Dispose();
                return;
            }

            RepositionOverlay();
        }

        private void OnWindowFound()
        {
            ApplyCaptionColor();
            CreateOverlay();
            RepositionOverlay();
        }

        private void ApplyCaptionColor()
        {
            // DWMWA_CAPTION_COLOR/TEXT_COLOR only take effect on Windows 11
            // (build 22000+). On older builds these calls fail silently and
            // the overlay tag remains the only visible identifier.
            try
            {
                int colorRef = ToColorRef(TagColor);
                DwmSetWindowAttribute(_rdpHwnd, DWMWA_CAPTION_COLOR, ref colorRef, sizeof(int));

                double luminance = (0.299 * TagColor.R + 0.587 * TagColor.G + 0.114 * TagColor.B) / 255.0;
                Color textColor = luminance > 0.6 ? Color.Black : Color.White;
                int textColorRef = ToColorRef(textColor);
                DwmSetWindowAttribute(_rdpHwnd, DWMWA_TEXT_COLOR, ref textColorRef, sizeof(int));
            }
            catch
            {
                // Ignore - unsupported OS version.
            }
        }

        private static int ToColorRef(Color c) => c.R | (c.G << 8) | (c.B << 16); // COLORREF is 0x00BBGGRR

        private void CreateOverlay()
        {
            _overlay = new TagOverlayForm(Label, TagColor);
            _overlay.Show();

            int exStyle = GetWindowLong(_overlay.Handle, GWL_EXSTYLE);
            exStyle |= WS_EX_LAYERED | WS_EX_TOOLWINDOW | WS_EX_NOACTIVATE;
            SetWindowLong(_overlay.Handle, GWL_EXSTYLE, exStyle);
        }

        private void RepositionOverlay()
        {
            if (_overlay == null || _overlay.IsDisposed)
            {
                return;
            }

            if (!GetWindowRect(_rdpHwnd, out RECT rect))
            {
                return;
            }

            // Hide the tag while the RDP window is minimized; there's no
            // caption bar to sit on in that state.
            if (IsIconic(_rdpHwnd))
            {
                if (_overlay.Visible)
                {
                    _overlay.Hide();
                }

                return;
            }
            if (!_overlay.Visible)
            {
                _overlay.Show();
            }

            // Sit the tag near the left of the caption bar, just right of
            // the system icon.
            int tagX = rect.Left + (rect.Right - rect.Left) / 2 + 46;
            int tagY = rect.Top + 6;

            SetWindowPos(_overlay.Handle, HWND_TOPMOST, tagX, tagY,
                _overlay.Width, _overlay.Height, SWP_NOACTIVATE | SWP_SHOWWINDOW);
        }

        private static IntPtr FindMstscWindowForProcess(int processId)
        {
            IntPtr found = IntPtr.Zero;
            EnumWindows((hwnd, _) =>
            {
                GetWindowThreadProcessId(hwnd, out int pid);
                if (pid != processId || !IsWindowVisible(hwnd))
                {
                    return true;
                }

                var sb = new StringBuilder(256);
                GetClassName(hwnd, sb, sb.Capacity);
                if (sb.ToString() == "TscShellContainerClass")
                {
                    found = hwnd;
                    return false; // stop enumeration
                }
                return true;
            }, IntPtr.Zero);
            return found;
        }

        public void Dispose()
        {
            if (_disposed)
            {
                return;
            }

            _disposed = true;
            _watchTimer.Stop();
            _watchTimer.Dispose();
            _overlay?.Close();
            _overlay?.Dispose();
        }
    }
}