import struct
import zlib
import re
import io
import os

from dataclasses import dataclass

from Crypto.Cipher import AES


SOURCE = "./kcf"
OUT = "templates/pages/sega/software/security/"

APPDATA_BINARY_LEN = 288  # RingEdge
APPDATA_BINARY_LEN2 = 416  # RingEdge 2
APPDATA_NU_BINARY_LEN = 112  # Nu

FUNCTION_TYPE = {
    1: "Server (SV)",
    2: "Satellite (ST)",
    3: "Live (LV)",
    4: "Terminal (??)",
    5: "UNKNOWN:5",
    11: "UNKNOWN:11",
    12: "UNKNOWN:12",
    21: "UNKNOWN:21",
    22: "UNKNOWN:22",
    98: "Standalone (??)",
}
PLATFORM = {
    "AAM": "RingWide",
    "AAL": "RingEdge",
    "AAS": "RingEdge 2",
    "AAR": "ELEFUN",

    "AAV": "Nu",
    "AAZ": "Nu Lv 3.1",
    "AAW": "Nu SX",
    "ACA": "ALLS X",
    "ACB": "ALLS X2",
}

SF_DEVELOP = 1 << 0
SF_ALL_NET = 1 << 2
SF_DELIVER = 1 << 3
SF_BINDING = 1 << 4
SF_BILLING = 1 << 5
SF_RENTAL = 1 << 6


class Stream:
    def __init__(self, stream):
        self.stream = stream

    def read(self, n):
        return self.stream.read(n)

    def seek(self, n):
        self.stream.seek(n)

    def crc32(self, nbytes):
        start = self.stream.tell()
        data = self.stream.read(nbytes)
        self.stream.seek(start)
        return zlib.crc32(data) & 0xffffffff

    def skip(self, n):
        self.read(n)

    def unpack(self, fmt):
        return struct.unpack("<" + fmt, self.stream.read(struct.calcsize(fmt)))

    def u8(self):
        return self.unpack("B")[0]

    def u16(self):
        return self.unpack("H")[0]

    def u32(self):
        return self.unpack("I")[0]

    def ipv4(self):
        return self.unpack("4B")

    def ipv6(self):
        return self.unpack("8H")

    def str(self, n):
        return self.unpack(f"{n}s")[0].decode("latin-1")


class AppData:
    system_flag: int
    function_type: int
    region: int
    platform_id: str

    def parse_region(self):
        if self.region == 0xFF:
            return "ALL"
        regions = []
        if self.region & 1:
            regions.append("JPN")
        if self.region & 2:
            regions.append("USA")
        if self.region & 4:
            regions.append("EXP")
        if self.region & 8:
            regions.append("CHN")
        if self.region & ~(1 | 2 | 4 | 8):
            regions.append(f"ex:{self.region}")
        return "/".join(regions)

    def parse_system_flag(self):
        """
        [bit 0] Develop:
            1: Key chip for development
            0: Key chip for mass production

        [bit 2] ALL.Net:
            1: ON (use)
            0: OFF (do not use)

        [bit 3] Deliver (LAN installation):
            1: Local FDC Server
            0: Client
            Ethernet network: Used & Mode: 1 only when server

        [bit 4] Binding:
            1: ON
            0: OFF
            * Not used in ELEFUN (= 0) (cannot be set as = 1)

        [bit 5] Billing:
            Requires ALL.Net to be ON
            1: ON
            0: OFF (no charge)
            * Not used in ELEFUN (= 0) (cannot be set as = 1)

        [bit 6] Rental (P-ras charging):
            Requires ALL.Net to be ON
            Requires Billing to be ON
            * Not used in ELEFUN (= 0) (cannot be set as = 1)

            1: P-ras billing (rental billing)
            0: pay-as-you-go billing (net billing)
        """

        parsed = []
        if self.system_flag & 1:
            parsed.append("Develop")
        if self.system_flag & 2:
            parsed.append("2")
        if self.system_flag & 4:
            parsed.append("ALL.Net")
        if self.system_flag & 8:
            parsed.append("Deliver")
        if self.system_flag & 16:
            parsed.append("Binding")
        if self.system_flag & 32:
            parsed.append("Billing")
        if self.system_flag & 64:
            parsed.append("Rental")
        if self.system_flag & 128:
            parsed.append("128")
        return ",".join(parsed)

    def parse_function_type(self):
        return FUNCTION_TYPE[self.function_type]

    def parse_platform_id(self):
        return f"{PLATFORM[self.platform_id]} ({self.platform_id})"


@dataclass
class AppDataRing(AppData):
    # Appboot
    format_type: int
    game_id: str
    region: int
    function_type: int
    system_flag: int
    platform_id: str
    dvd_flag: int
    network_addr: tuple[int]

    # Misc
    app_data: bytes

    # Crypto
    key: bytes
    iv: bytes
    seed: bytes
    keyfile: bytes

    # Misc (2)
    extra: bytes


@dataclass
class AppDataNu(AppData):
    # Appboot
    format_type: int
    game_id: str
    region: int
    function_type: int
    system_flag: int
    billing_form: int
    platform_id: str
    network_addr: tuple[int]
    network_addr_6: tuple[int]

    # Crypto
    gkey: bytes
    giv: bytes
    okey: bytes
    oiv: bytes


def scramble(val, a, b, c, d):
    ret = bytearray(val)
    ret[a], ret[b] = ret[b], ret[a]
    ret[c], ret[d] = ret[d], ret[c]
    return bytes(ret)


def parse_appdata_ring(stream: Stream):
    crc = stream.u32()  # idk how to make this pass. it doesn't match?

    format_type = stream.u8()
    stream.skip(3)
    game_id = stream.str(4)
    region = stream.u8()
    function_type = stream.u8()
    system_flag = stream.u8()
    stream.skip(1)
    platform_id = stream.str(3)
    dvd_flag = stream.u8()
    network_addr = stream.ipv4()
    app_data = stream.read(216)
    seed = scramble(stream.read(16), 1, 8, 12, 15)
    key = scramble(stream.read(16), 0, 4, 2, 14)
    iv = scramble(stream.read(16), 0, 11, 5, 15)

    # AAL, AAS, AAR has:
    extra = stream.read(128)
    # For AAR, extra=FF....FF

    # if extra and platform_id != 'AAR':
    #     print(extra[:64].hex(), platform_id, game_id, function_type)

    cipher = AES.new(key, AES.MODE_CBC, iv)
    keyfile = cipher.encrypt(seed)

    return AppDataRing(
        format_type, game_id, region, function_type, system_flag,
        platform_id, dvd_flag, network_addr, app_data,
        key, iv, seed, keyfile, extra
    )


def parse_appdata_nu(stream: Stream):
    crc = stream.u32()

    format_type = stream.u8()
    stream.skip(3)
    game_id = stream.str(4)
    region = stream.u8()
    function_type = stream.u8()
    system_flag = stream.u8()
    billing_form = stream.u8()
    platform_id = stream.str(3)
    stream.skip(1)
    network_addr = stream.ipv4()
    stream.skip(8)
    ipv6_addr = stream.ipv6()
    gkey = stream.read(16)
    giv = stream.read(16)
    okey = stream.read(16)
    oiv = stream.read(16)

    return AppDataNu(
        format_type, game_id, region, function_type, system_flag,
        billing_form, platform_id, network_addr, ipv6_addr,
        gkey, giv, okey, oiv
    )


def format_row(*values, elem="td"):
    return "<tr>" + "".join(f"<{elem}>{i}</{elem}>" for i in values) + "</tr>"


def format_code(code):
    return f"<code>{code}</code>"


def parse_kcf(data):
    kcf = io.BytesIO(data)

    nbytes = len(data)

    if nbytes == APPDATA_BINARY_LEN:
        return parse_appdata_ring(Stream(kcf)), None
    elif nbytes == APPDATA_BINARY_LEN2:
        return parse_appdata_ring(Stream(kcf)), None
    elif nbytes == APPDATA_NU_BINARY_LEN:
        return None, parse_appdata_nu(Stream(kcf))
    else:
        raise ValueError


STYLES = """
<style>
body {
    font-family: sans-serif;
}
table {
    white-space: nowrap;
    overflow-x: auto;
    border-collapse: collapse;
}
th {
    border-bottom: 3px solid #000;
    position: sticky;
    top: 0;
    text-align: left;
    background: #fff;
}
td {
    border-bottom: 1px solid #ddd;
}
th, td {
    border-left: 1px solid #ddd;
}
th:last-child, td:last-child {
    border-right: 1px solid #ddd;
}
th, td {
    margin: 0;
    padding: 2px 4px;
}
</style>
"""


def gen_file_ring(path: str, platform: str, appdata: list[AppDataRing]):
    with open(path, "w") as kcf_html:
        print(STYLES, file=kcf_html)
        print("<table><thead>", file=kcf_html)
        print(format_row(
            "Game ID", "Region", "Model Type",
            "Develop Flag", "ALL.Net", "Deliver", "Binding", "Billing", "Rental",
            "DVD Enabled", "Platform ID", "IPv4 Subnet",
            "AES Key", "AES IV", "AppBoot Seed", "Keyfile",
            "App Data", "Trailing Data",
            elem="th"
        ), file=kcf_html)
        print("</thead></tbody>", file=kcf_html)
        appdata.sort(key=lambda x: (x.platform_id, x.network_addr, x.game_id, x.format_type, x.system_flag))

        for i in appdata:
            if i.platform_id != platform:
                continue

            print(format_row(
                i.game_id,
                i.parse_region(),
                i.parse_function_type(),
                "X" if i.system_flag & SF_DEVELOP else "",
                "X" if i.system_flag & SF_ALL_NET else "",
                "X" if i.system_flag & SF_DELIVER else "",
                "X" if i.system_flag & SF_BINDING else "",
                "X" if i.system_flag & SF_BILLING else "",
                "X" if i.system_flag & SF_RENTAL else "",
                "X" if i.dvd_flag else "",
                i.parse_platform_id(),
                ".".join(map(str, i.network_addr)),
                format_code(i.key.hex().upper()),
                format_code(i.iv.hex().upper()),
                format_code(i.seed.hex().upper()),
                format_code(i.keyfile.hex().upper()),
                ""
                if all(j == 0 for j in i.app_data)
                else format_code(i.app_data.hex().upper()),
                format_code(i.extra.hex().upper()),
            ), file=kcf_html)
        print("</tbody></table>", file=kcf_html)


def main():
    appdata_ring: list[AppDataRing] = []
    appdata_nu: list[AppDataNu] = []

    raw = set()
    for i in os.listdir(SOURCE):
        raw.add(open(os.path.join(SOURCE, i), "rb").read())

    for i in raw:
        ring, nu = parse_kcf(i)

        if ring:
            appdata_ring.append(ring)
        if nu:
            appdata_nu.append(nu)

    gen_file_ring(OUT + "kcf-ringedge.html", "AAL", appdata_ring)
    gen_file_ring(OUT + "kcf-ringwide.html", "AAM", appdata_ring)
    gen_file_ring(OUT + "kcf-ringedge2.html", "AAS", appdata_ring)
    gen_file_ring(OUT + "kcf-elefun.html", "AAR", appdata_ring)

    with open(OUT + "kcf-nu.html", "w") as kcf_nu:
        print(STYLES, file=kcf_nu)
        print("<table><thead>", file=kcf_nu)
        print(format_row(
            "Game ID", "Region", "Model Type",
            "Develop Flag", "ALL.Net", "Deliver", "Binding", "Billing", "Rental",
            "Billing Form", "Platform ID",
            "IPv4 Subnet", "IPv6 Subnet",
            "Game Key", "Game IV", "Opt Key", "Opt IV",
            elem="th"
        ), file=kcf_nu)
        print("</thead></tbody>", file=kcf_nu)
        appdata_nu.sort(key=lambda x: (x.platform_id, x.network_addr, x.game_id, x.format_type, x.system_flag))

        for i in appdata_nu:
            ipv6 = ":".join(f"{j:04x}" for j in i.network_addr_6)
            ipv6 = re.sub(r"/(^(0000:?)+)|(:(0000:)+)|((:0000)+$)/", "::", ipv6)

            print(format_row(
                i.game_id,
                i.parse_region(),
                i.parse_function_type(),
                "X" if i.system_flag & SF_DEVELOP else "",
                "X" if i.system_flag & SF_ALL_NET else "",
                "X" if i.system_flag & SF_DELIVER else "",
                "X" if i.system_flag & SF_BINDING else "",
                "X" if i.system_flag & SF_BILLING else "",
                "X" if i.system_flag & SF_RENTAL else "",
                i.billing_form,
                i.parse_platform_id(),
                ".".join(map(str, i.network_addr)),
                ipv6,
                format_code(i.gkey.hex().upper()),
                format_code(i.giv.hex().upper()),
                format_code(i.okey.hex().upper()),
                format_code(i.oiv.hex().upper()),
            ), file=kcf_nu)
        print("</tbody></table>", file=kcf_nu)


if __name__ == "__main__":
    main()