#!/usr/bin/env python3

template = """
#-------------------------------------------------------------
#
#  License created by Gowin_LIC_Generator
#
# WARNING: do NOT modify any part of this file.
#
#-------------------------------------------------------------

MODE = NODELOCK
TYPE = NODE_LOCK
HOST_ID = {host_id}
EXP_DATE = 2100-01-01
TOOL = {{gowin, 2100-01-01, 0}}
""".lstrip()

begin_sign = b"""
-------------------------- BEGIN SIGN --------------------------
""".strip()
end_sign = b"""
--------------------------- END SIGN ---------------------------
""".strip()

dll_template = """
#define _GNU_SOURCE
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <dlfcn.h>

const char orig[] = {orig};
const char patched[] = {patched};

static void *(*real_BIO_new_mem_buf)(const void *buf, int len) = NULL;

void init() __attribute__((constructor));
void init() {{
  real_BIO_new_mem_buf = dlsym(RTLD_NEXT, "BIO_new_mem_buf");
}}
void *BIO_new_mem_buf(const void *buf, int len) {{
  if (strcmp(orig, buf) == 0) {{
    return real_BIO_new_mem_buf(patched, len);
  }}
  return real_BIO_new_mem_buf(buf, len);
}}
"""

orig = """
-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDicOsV7qRM8ReSKExxF0wD+xF7
Jc1YqXFejXu9XpONB+yjZ4X+KUkYL2AY6AWFDlV0KFWccG/CCGrJRetLrg/EqcbU
7laQnU5wJSNlWf+n0g93g4JHVm1KClpJUXIq1Jb9WGR3mG6rmKRGp+iktqbn5Gwh
sDdmPxJC+BWJj73BpwIDAQAB
-----END PUBLIC KEY-----
""".strip()

def run():
    import Crypto
    from Crypto.PublicKey import RSA
    from Crypto.Hash import SHA256
    from Crypto.Signature import PKCS1_v1_5
    import uuid

    import argparse
    import sys, os
    import subprocess, tempfile
    import logging

    p = argparse.ArgumentParser("GoWin license generator")
    p.add_argument("-v", "--verbose",
                   action = "store_true",
                   default = False)
    p.add_argument("--host-id",
                   default = f"{uuid.getnode():012x}",
                   type = str,
                   help = f"eth0 MAC addr, default = {uuid.getnode():012x}")
    p.add_argument("--private-key",
                   type = argparse.FileType("rb"),
                   nargs = "?",
                   help = "Private key")
    p.add_argument("--write-private-key",
                   type = argparse.FileType("wb"),
                   nargs = "?",
                   help = "Where to put generated private key")
    p.add_argument("--dll",
                   type = argparse.FileType("rb"),
                   nargs = "?",
                   help = "Dll file location")
    p.add_argument("dest",
                   type = argparse.FileType("wb"),
                   nargs = "?",
                   default = sys.stdout,
                   help = "Destination file location (default = stdout)")
    r = p.parse_args(sys.argv[1:])

    if r.verbose:
        logging.basicConfig(level = logging.DEBUG)
    else:
        logging.basicConfig(level = logging.INFO)
    msg = template.format(host_id = r.host_id).encode()

    if r.private_key:
        logging.info(f"Using private key from \"{r.private_key.name}\"")
        key = RSA.import_key(r.private_key.read())
    else:
        logging.info(f"Generating new private key")
        key = RSA.generate(1024)

    if r.write_private_key:
        r.write_private_key.write(key.export_key())

    for e in key.export_key().decode().split("\n"):
        logging.debug(f"\t{e}")
    for e in key.public_key().export_key().decode().split("\n"):
        logging.debug(f"\t{e}")

    digest = SHA256.new(msg)
    signature = PKCS1_v1_5.new(key).sign(digest)
    signature_text = b"\n".join([signature[i:32+i].hex().encode() for i in range(0, len(signature), 32)])

    os.write(r.dest.fileno(), msg)
    lines = [begin_sign, signature_text, end_sign, b""]
    os.write(r.dest.fileno(), b"\n".join(lines))

    if r.dll:
        logging.info(f"Compiling {r.dll.name}")
        with tempfile.NamedTemporaryFile(mode = "w", prefix = "gowin_dll_", suffix = ".c") as fd:
            patched = key.public_key().export_key().decode()
            content = dll_template.format(
                orig = '"' + "\\n".join(orig.split("\n")) + '"',
                patched = '"' + "\\n".join(patched.split("\n")) + '"')
            fd.write(content)

            fd.flush()
            subprocess.run(["gcc", "-Werror", "-Wall", "-O2", "-shared", "-fPIC", fd.name, "-o", r.dll.name],
                           stderr = subprocess.STDOUT)

    logging.info("Done")
    logging.info(f"Use \"{r.dest.name}\" as license file")
    if r.dll:
        logging.info(f"Run GoWin software using LD_PRELOAD=\"<...>/{r.dll.name}\"")
    else:
        orig_sig = RSA.import_key(orig).export_key().hex()
        rep_sig = key.public_key().export_key().hex()
        logging.info(f"Use sfk replace -yes -binary /{orig_sig}/{rep_sig}/ -dir <...>")

if __name__ == "__main__":
    run()