#!/usr/bin/env python3

# Pollute Proxmox VE 5.x subscription cache so it won't alert you on dashboard login
# If you need to prevent it checking keys against a server, please block "shop.maurer-it.com"
import hashlib
import base64
import json
import time
import re
import sys
import os
from typing import List

# PVE & PMG: /usr/share/perl5/PVE/Subscription.pm
# PBS: /usr/lib/x86_64-linux-gnu/proxmox-backup/*
shared_key_data = "kjfdlskfhiuewhfk947368"
server_key_file = "/etc/ssh/ssh_host_rsa_key.pub"

def get_timestamp() -> int:
    return int(time.time())

# perl's md5_base64 implementation
def md5_base64(x: str) -> str: 
    return base64.b64encode(hashlib.md5(x.encode()).digest()).strip(b'=')

def generate_server_id(key: str) -> str:
    return hashlib.md5(key.encode()).hexdigest().upper()

def generate_subscription(key: str, server_ids: List[str]) -> str:
    localinfo = {
        "checktime": get_timestamp(),
        "status": "Active", # PBS: `new`, `notfound`, `active`, `invalid`
        "key": key,
        "validdirectory": ",".join(server_ids),
        "productname": "YajuuSenpai",
        "regdate": get_timestamp(),
        "nextduedate": 2147483647,
    }

    data = base64.standard_b64encode(json.dumps(localinfo).encode()).decode()
    cat = str(localinfo["checktime"]) + data + "\n" + shared_key_data
    csum = md5_base64(cat).decode()

    return key + "\n" + csum + "\n" + data + "\n"

# key_pattern can be find in /usr/share/perl5/{PVE,PMG}/API2/Subscription.pm
def activate(key: str, key_pattern: str, subscription_file: str) -> None:
    # check if the key format is correct
    pattern = re.compile(key_pattern)
    if not pattern.match(key):
        print("key format error", file=sys.stderr)
        sys.exit(1)

    # get machine ID
    server_id = ""
    with open(server_key_file, "r") as f:
        server_id = generate_server_id(f.read())

    # generate a license file
    subscription = generate_subscription(key, [server_id])

    # write license file
    with open(subscription_file, "w") as f:
        f.write(subscription)

if __name__ == "__main__":
    # Proxmox VE
    if os.path.exists("/etc/pve"):
        print("Activating Proxmox VE...")
        activate("pve8p-1145141919", r'pve([1248])([cbsp])-[0-9a-f]{10}', "/etc/subscription")

    # Proxmox Mail Gateway
    if os.path.exists("/etc/pmg"):
        print("Activating Proxmox Mail Gateway...")
        activate("pmgp-1145141919", r'pmg([cbsp])-[0-9a-f]{10}', "/etc/pmg/subscription")
    
    # Proxmox Backup Server (not working yet)
    if os.path.exists("/etc/proxmox-backup"):
        print("Activating Proxmox Backup Server...")
        activate("pmgp-1145141919", r'pmg([cbsp])-[0-9a-f]{10}', "/etc/proxmox-backup/subscription")