bemaniutils/bemani/utils/responsegen.py

import argparse
import sys
from typing import Dict, List, Optional

from bemani.protocol import EAmuseProtocol, Node


def generate_node_name(node: Node, used_names: Dict[str, Node]) -> str:
    potential_name = node.name
    if potential_name not in used_names:
        used_names[potential_name] = node
        return potential_name

    loop = 1
    while f'{potential_name}_{loop}' in used_names:
        loop = loop + 1

    potential_name = f'{potential_name}_{loop}'
    used_names[potential_name] = node
    return potential_name


def generate_node_create(node: Node) -> str:
    dtype = node.data_type
    if dtype == 'str':
        method = 'string'
    elif dtype == 'bin':
        method = 'binary'
    elif dtype == 'ip4':
        method = 'ipv4'
    elif dtype == '4u8':
        method = 'fouru8'
    else:
        method = dtype
    if node.is_array:
        method = f'{method}_array'

    if dtype != 'void':
        # Format the type for display
        if dtype == 'str':
            value = f', \'{node.value}\''
        elif dtype == 'ip4':
            value = f', \'{node.value[0]}.{node.value[1]}.{node.value[2]}.{node.value[3]}\''
        else:
            value = f', {node.value}'
    else:
        value = ''

    return f'Node.{method}(\'{node.name}\'{value})'


def generate_node_link(node_name: str, used_names: Dict[str, Node], parent: Node) -> str:
    # Find the node that parents this, link to it
    found_parent = None
    for parent_name in used_names:
        if used_names[parent_name] is parent:
            found_parent = parent_name
            break

    if found_parent is None:
        raise Exception(f'Failed to find parent name for {parent}')

    return f'{found_parent}.add_child({node_name})'


def generate_lines(node: Node, used_names: Dict[str, Node], parent: Optional[Node]=None) -> List[str]:
    # First, generate node itself
    create = generate_node_create(node)
    if not node.children and not node.attributes and parent:
        # Just directly hook this up to parent
        return [generate_node_link(create, used_names, parent)]

    # Print the node generate itself
    out = []
    node_name = generate_node_name(node, used_names)
    out.append(f'{node_name} = {create}')

    # Now, generate add to parent if exists
    if parent is not None:
        out.append(generate_node_link(node_name, used_names, parent))

    # Now generate node attributes
    for attr in node.attributes:
        out.append(f'{node_name}.set_attribute(\'{attr}\', \'{node.attributes[attr]}\')')

    # Now, do the same for all children
    for child in node.children:
        out.extend(generate_lines(child, used_names, node))

    return out


def generate_code(infile: str, outfile: str, encoding: str) -> None:
    if infile == '-':
        # Load from stdin
        packet = sys.stdin.buffer.read()
    else:
        with open(infile, mode="rb") as infp:
            packet = infp.read()
            infp.close

    # Add an XML special node to force encoding (will be overwritten if there
    # is one in the packet).
    packet = b''.join([
        f'<?xml encoding="{encoding}"?>'.encode(encoding),
        packet,
    ])

    # Attempt to decode it
    proto = EAmuseProtocol()
    req = proto.decode(
        None,
        None,
        packet,
    )

    if req is None:
        # Can't decode, exit
        raise Exception("Unable to decode packet!")

    # Walk through, outputting each node and attaching it to its parent
    code = '\n'.join(generate_lines(req, {}))

    if outfile == '-':
        print(code)
    else:
        with open(outfile, mode="a") as outfp:
            outfp.write(code)
            outfp.close


def main() -> None:
    parser = argparse.ArgumentParser(description="A utility to generate code that will generate a packet.")
    parser.add_argument("-i", "--infile", help="File containing an XML or binary node structure. Use - for stdin.", type=str, default=None, required=True)
    parser.add_argument("-o", "--outfile", help="File to write python code to. Use - for stdout.", type=str, default=None, required=True)
    parser.add_argument("-e", "--encoding", help="Encoding for the packet, defaults to UTF-8.", type=str, default='utf-8')
    args = parser.parse_args()

    generate_code(args.infile, args.outfile, args.encoding)


if __name__ == '__main__':
    main()