#!/usr/bin/env python3

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#

# This script prints some information about an executable's sections to stdout,
# depending on the command line options passed:
# * With --text-size: the program outputs the size of the `.text` segment (or
#   equivalent) as an integer (in bytes)
# * Without --text-size: a JSON is returned containing information about all
#   sections

import argparse
import json
import sys
from subprocess import DEVNULL, PIPE, run

import yaml


def find_section_attribute(data: dict, name: str, attr_name: str) -> int | None:
    for section in data["Sections"]:
        if section["Name"]["Value"] == name:
            return section[attr_name]
    return None


def _set_or_add(num1: int | None, num2: int):
    return num2 if num1 is None else num1 + num2


def _flag_overlap(flag_set: dict, flags: frozenset[str]):
    new_flags = {f["Name"] for f in flag_set}
    return bool(new_flags & flags)


def compute_segment_load_size(data: dict):
    size = None
    for segment in data["ProgramHeaders"]:
        if segment["Type"]["Value"] != "PT_LOAD":
            continue
        if _flag_overlap(segment["Flags"]["Flags"], frozenset(("PF_X",))):
            size = _set_or_add(size, segment["FileSize"])
    return size


def compute_executable_sections_macho(data: dict):
    size = None
    executable_flags = frozenset(("PureInstructions", "SomeInstructions"))
    for section in data["Sections"]:
        if _flag_overlap(section["Attributes"]["Flags"], executable_flags):
            size = _set_or_add(size, section["Size"])
    return size


def compute_executable_sections_pe(data: dict):
    size = None
    executable_flags = frozenset(("IMAGE_SCN_CNT_CODE", "IMAGE_SCN_MEM_EXECUTE"))
    for section in data["Sections"]:
        if _flag_overlap(section["Characteristics"]["Flags"], executable_flags):
            size = _set_or_add(size, section["RawDataSize"])
    return size


def get_text_size(data: dict):
    format_: str = data["Format"]
    if format_.startswith("elf"):
        res = find_section_attribute(data, ".text", "Size")
        if res is None:
            res = compute_segment_load_size(data)
    elif format_.startswith("Mach-O"):
        res = find_section_attribute(data, "__text", "Size")
        if res is None:
            res = compute_executable_sections_macho(data)
    elif format_.startswith("COFF"):
        res = find_section_attribute(data, ".text", "RawDataSize")
        if res is None:
            res = compute_executable_sections_pe(data)
    else:
        raise ValueError(f"Invalid format: {format_}")

    return res if res is not None else -1


def readobj(filepath: str):
    proc = run(
        ["llvm-readobj", "--elf-output-style=JSON", "--sections", "--segments", filepath],
        stdout=PIPE,
        stderr=DEVNULL,
        check=True,
        text=True,
    )
    data = yaml.safe_load(proc.stdout)

    result = {}
    if len(data) == 1:
        # ELF output
        result.update(data[0]["FileSummary"])
        result["Sections"] = []
        for section in data[0]["Sections"]:
            result["Sections"].append(section["Section"])
        result["ProgramHeaders"] = []
        for segment in data[0]["ProgramHeaders"]:
            result["ProgramHeaders"].append(segment["ProgramHeader"])
    elif len(data) == 5:
        # PE/MACH-O output
        for i in range(0, 4):
            result.update(data[i])
        result["Sections"] = []
        for section in data[4]["Sections"]:
            result["Sections"].append(section["Section"])
    else:
        raise ValueError()
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--text-size", action="store_true", help="Output the size of the .text section"
    )
    parser.add_argument("input", help="Input file")
    args = parser.parse_args()

    readobj_data = readobj(args.input)
    if args.text_size:
        print(get_text_size(readobj_data))
    else:
        json.dump(readobj_data, sys.stdout)


if __name__ == "__main__":
    main()