warp_tab_contrast_patch.py

#!/usr/bin/env python3
"""
Warp Terminal - Active Tab Contrast Patch
=========================================

Patches the WGSL rect shader embedded in warp.exe to increase contrast
between active and inactive tabs in dark mode.

Problem:
    Warp's dark mode renders the active tab only ~4% brighter than inactive
    tabs (#121212 vs #1D1D1D with black background), making it nearly
    impossible to tell which tab is selected.

Solution:
    Modifies the GPU rect fragment shader to invert and amplify the
    brightness relationship in the tab bar area (top 80px of the window):
    - Active tabs (originally darkest) become visibly bright (~#404040)
    - Inactive tabs (originally slightly lighter) become very dark (~#0D0D0D)

    The patch only affects dark surfaces (luminance < 0.15) in the top 80
    pixels of the viewport, so the terminal content area is unaffected.

Usage:
    1. Close Warp completely
    2. Run: python warp_tab_contrast_patch.py
    3. Relaunch Warp

    To restore the original binary:
        python warp_tab_contrast_patch.py --restore

    To adjust the effect strength:
        python warp_tab_contrast_patch.py --strength 0.8   (stronger)
        python warp_tab_contrast_patch.py --strength 0.4   (subtler)

    Default strength is 0.6. The formula applied in the shader is:
        adjustment = strength - luminance * (strength * 10.0)

Tested with Warp for Windows. The script auto-detects the shader location
in the binary, so it should survive Warp updates as long as the shader
source text remains similar.

Backup:
    A backup (warp.exe.bak) is created automatically before the first patch.
    Use --restore to revert to it.
"""

import argparse
import os
import shutil
import struct
import sys

# Default Warp installation path on Windows
DEFAULT_WARP_DIR = os.path.join(
    os.environ.get("LOCALAPPDATA", ""), "Programs", "Warp"
)
DEFAULT_WARP_EXE = os.path.join(DEFAULT_WARP_DIR, "warp.exe")
BACKUP_SUFFIX = ".bak"

# Shader identification markers
SHADER_START_MARKER = b"Rect instance bufferRect Shader"
SHADER_RETURN_STMT = "return background_color;"

# Comments to remove to make room for the patch (order matters: largest first)
COMMENTS_TO_REMOVE = [
    # In the Uniforms struct
    (
        "// Padding necessary to ensure that the uniforms is 16 bytes."
        " Some wgpu-supported devices (such as webgl) require\r\n"
        "    // buffer bindings to be a multiple of 16 bytes.\r\n"
    ),
    # In the vertex shader input struct
    (
        "// The sigma and padding factor values packed into a single vec2."
        " We pack them together in order\r\n"
        "    // to reduce the total number of attributes, which maxes out"
        " at 16. See here:\r\n"
        "    // https://docs.rs/wgpu/latest/wgpu/struct.Limits.html"
        "#structfield.max_vertex_attributes\r\n"
    ),
    # Dashed border data comment
    (
        "// The length of the dash and the gaps for the x and y"
        " dimensions, packed into a single vec3.\r\n"
    ),
]

# Patch marker to detect already-patched binaries
PATCH_MARKER = "/*TCP*/"  # Tab Contrast Patch


def find_shader(data: bytes) -> tuple[int, int]:
    """Find the rect shader string in the binary. Returns (start, end) offsets."""
    start = data.find(SHADER_START_MARKER)
    if start == -1:
        return -1, -1

    # Find the null terminator marking the end of the shader string
    end = start
    while end < len(data) and data[end] != 0:
        end += 1

    return start, end


def build_contrast_code(strength: float) -> str:
    """Build the WGSL contrast injection code.

    The code computes the luminance of each dark rect in the tab bar area
    and applies a linear remap that brightens the darkest surfaces (active
    tabs) and darkens the slightly-lighter surfaces (inactive tabs).

    Args:
        strength: Controls the effect intensity. Default 0.6.
                  Higher = more contrast. Range roughly [0.2, 1.0].
    """
    slope = strength * 10.0
    return (
        f"{PATCH_MARKER}"
        f"let cl=dot(background_color.rgb,vec3<f32>(0.299,0.587,0.114));"
        f"if cl>0.005&&cl<0.15&&in.position.y<80.0{{"
        f"let a={strength}-cl*{slope};"
        f"background_color=vec4<f32>("
        f"background_color.rgb+vec3<f32>(a,a,a),background_color.a);"
        f"}}\r\n    "
    )


def patch_shader(shader: str, strength: float) -> str:
    """Apply the contrast patch to the shader source.

    Removes select comments to make room, then inserts the contrast code
    just before the fragment shader's return statement. The result is
    padded with spaces to maintain the exact original byte count.
    """
    original_len = len(shader)

    if PATCH_MARKER in shader:
        raise RuntimeError(
            "Binary is already patched. Use --restore first, then re-patch."
        )

    if SHADER_RETURN_STMT not in shader:
        raise RuntimeError(
            f"Could not find '{SHADER_RETURN_STMT}' in the rect shader. "
            "The shader format may have changed in this Warp version."
        )

    # Remove comments to free space
    for comment in COMMENTS_TO_REMOVE:
        if comment in shader:
            shader = shader.replace(comment, "", 1)

    # Insert contrast code before the return statement
    contrast_code = build_contrast_code(strength)
    shader = shader.replace(
        SHADER_RETURN_STMT,
        contrast_code + SHADER_RETURN_STMT,
        1,
    )

    # Pad with spaces to maintain exact byte count
    diff = original_len - len(shader)
    if diff < 0:
        raise RuntimeError(
            f"Patch is {-diff} bytes too long. Not enough comments to remove. "
            "Try a lower --strength value or report this issue."
        )

    if diff > 0:
        # Insert padding spaces before the return statement
        shader = shader.replace(
            contrast_code + SHADER_RETURN_STMT,
            contrast_code + (" " * diff) + SHADER_RETURN_STMT,
            1,
        )

    assert len(shader) == original_len, (
        f"Internal error: size mismatch ({len(shader)} vs {original_len})"
    )
    return shader


def main():
    parser = argparse.ArgumentParser(
        description="Patch Warp terminal for better active tab contrast in dark mode."
    )
    parser.add_argument(
        "--warp-exe",
        default=DEFAULT_WARP_EXE,
        help=f"Path to warp.exe (default: {DEFAULT_WARP_EXE})",
    )
    parser.add_argument(
        "--restore",
        action="store_true",
        help="Restore the original warp.exe from backup",
    )
    parser.add_argument(
        "--strength",
        type=float,
        default=0.6,
        help="Effect strength (default: 0.6). Higher = more contrast. Range: 0.2-1.0",
    )
    args = parser.parse_args()

    warp_exe = args.warp_exe
    backup_exe = warp_exe + BACKUP_SUFFIX

    if not os.path.isfile(warp_exe):
        print(f"Error: warp.exe not found at {warp_exe}", file=sys.stderr)
        print("Use --warp-exe to specify the correct path.", file=sys.stderr)
        sys.exit(1)

    # --- Restore mode ---
    if args.restore:
        if not os.path.isfile(backup_exe):
            print(f"Error: No backup found at {backup_exe}", file=sys.stderr)
            sys.exit(1)
        shutil.copy2(backup_exe, warp_exe)
        print(f"Restored original binary from {backup_exe}")
        return

    # --- Patch mode ---
    strength = args.strength
    if not 0.1 <= strength <= 1.5:
        print("Warning: strength outside recommended range [0.2, 1.0]", file=sys.stderr)

    # Read binary
    with open(warp_exe, "rb") as f:
        data = bytearray(f.read())

    # Find the shader
    shader_start, shader_end = find_shader(data)
    if shader_start == -1:
        print(
            "Error: Could not find the rect shader in warp.exe.\n"
            "This Warp version may not be supported.",
            file=sys.stderr,
        )
        sys.exit(1)

    shader_text = data[shader_start:shader_end].decode("utf-8")
    print(f"Found rect shader at offset 0x{shader_start:X} ({shader_end - shader_start} bytes)")

    # Check if already patched
    if PATCH_MARKER in shader_text:
        print(
            "Binary is already patched. Run with --restore first to re-patch "
            "with different settings.",
            file=sys.stderr,
        )
        sys.exit(1)

    # Create backup (only if one doesn't already exist)
    if not os.path.isfile(backup_exe):
        shutil.copy2(warp_exe, backup_exe)
        print(f"Created backup at {backup_exe}")
    else:
        print(f"Backup already exists at {backup_exe}")

    # Patch the shader
    try:
        patched_shader = patch_shader(shader_text, strength)
    except RuntimeError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    # Write patched binary
    patched_bytes = patched_shader.encode("utf-8")
    assert len(patched_bytes) == shader_end - shader_start
    data[shader_start:shader_end] = patched_bytes

    with open(warp_exe, "wb") as f:
        f.write(data)

    print(f"Patched successfully (strength={strength})")
    print("Active tab will be brighter, inactive tabs will be darker.")
    print("Launch Warp to see the change. Use --restore to revert.")


if __name__ == "__main__":
    main()