174 lines
4.5 KiB
Python
174 lines
4.5 KiB
Python
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 given an example packet from a log or binary dump."
|
|
)
|
|
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()
|