#!/usr/bin/env python3

# Pollute Proxmox software subscription cache so it won't alert you on dashboard login
# Should be scheduled to run every few hours with a timer to prevent cache from expiring
# If you need to prevent it checking keys against a server, please block "shop.maurer-it.com" in your hosts file

from __future__ import print_function
import base64
import hashlib
import json
import platform
import time
# import re
# import sys
import os
from datetime import datetime, timedelta

# PVE & PMG: /usr/share/perl5/PVE/Subscription.pm
# PBS: /usr/lib/x86_64-linux-gnu/proxmox-backup/* (source code at `https://git.proxmox.com/git/proxmox-backup.git`)
shared_key_data = "kjfdlskfhiuewhfk947368"
server_key_file = "/etc/ssh/ssh_host_rsa_key.pub"

# Default license keys
# https://github.com/proxmox/pve-manager/blob/c3a8ade7aceaafa4e73c4b0c4d93a1ee66db7913/PVE/API2/Subscription.pm#L46
lic_pve = {
    "aarch64": "pve8p-arm-1145141919",
}.get(platform.machine(), "pve8p-1145141919")
lic_pmg = "pmgp-1145141919"
lic_pbs = "pbst-1145141919"

# UI customization
ui_product_name = "YajuuSenpai"
ui_message = "Yajuu Senpai has got your back"
ui_url = "https://github.com/Jamesits/pve-fake-subscription"

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

# Perl's md5_base64 implementation
def md5_base64_perl(x):
    return base64.b64encode(hashlib.md5(x.encode()).digest()).strip(b'=').decode()

# Rust's `base64::encode(tools::md5sum("something")?);`
def md5_base64_rs(x):
    return base64.b64encode(hashlib.md5(x.encode()).digest()).decode()

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

def dt_string(format, offset_secs=0):
    return (datetime.now() + timedelta(seconds=offset_secs)).strftime(format)

def generate_subscription_pve_pmg(key, server_ids, product_name=ui_product_name):
    localinfo = {
        "checktime": get_timestamp(),
        "status": "Active",
        "key": key,
        "validdirectory": ",".join(server_ids),
        "productname": product_name,
        "regdate": dt_string("%Y-%m-%d %H:%M:%S"),
        "nextduedate": dt_string("%Y-%m-%d", 1296000),
    }

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

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

# key_pattern can be find in /usr/share/perl5/{PVE,PMG}/API2/Subscription.pm
# PVE9+: r'pve([1248])([cbsp])-(arm-)?[0-9a-f]{10}'
# PVE5-8: r'pve([1248])([cbsp])-[0-9a-f]{10}'
# PVE3-4: r'pve([124])([cbsp])-[0-9a-f]{10}'
# PMG: r'pmg([cbsp])-[0-9a-f]{10}'
def activate_pve_pmg(key, subscription_file, *args, **kwargs):
    # 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_pve_pmg(key, [server_id], *args, **kwargs)

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

def generate_subscription_pbs(key, server_ids, product_name=ui_product_name, message=ui_message, url=ui_url):
    localinfo = {
        "status": "active", # PBS: `new`, `notfound`, `active`, `invalid`
        "serverid": ",".join(server_ids),
        "checktime": get_timestamp(),
        "key": key,
        "message": message,
        "productname": product_name,
        "regdate": dt_string("%Y-%m-%d %H:%M:%S"),
        "nextduedate": dt_string("%Y-%m-%d", 1296000), # 1296000: MAX_LOCAL_KEY_AGE in src/tools/subscription.rs
        "url": url,
    }

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

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

# Key pattern: pbst-xxxxxxxxxx
def activate_pbs(key, subscription_file, *args, **kwargs):
    # 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_pbs(key, [server_id], *args, **kwargs)

    # 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_pve_pmg(lic_pve, "/etc/subscription")

    # Proxmox Mail Gateway
    if os.path.exists("/etc/pmg"):
        print("Activating Proxmox Mail Gateway...")
        activate_pve_pmg(lic_pmg, "/etc/pmg/subscription")

    # Proxmox Backup Server
    if os.path.exists("/etc/proxmox-backup"):
        print("Activating Proxmox Backup Server...")
        activate_pbs(lic_pbs, "/etc/proxmox-backup/subscription")
