Crimson Forge build pipeline tasks

tasks.py

import shutil
from pathlib import Path

from invoke import Collection, task

SRC_DIR = Path("src")
BIN_DIR = Path("bin")
GRAPHML_DIR = Path("graphml")
GRAPHVIZ_DIR = Path("graphviz")

CRIMSON_FORGE_IMAGE = "public.ecr.aws/n5b4u6h0/zerosteiner/crimson-forge:latest"
CRIMSON_FORGE_LOCAL_IMAGE = "crimson-forge:latest"


def _container_runtime():
    for runtime in ("podman", "docker"):
        if shutil.which(runtime):
            return runtime
    raise RuntimeError("Neither podman nor docker was found on PATH")


def _resolve_image(ctx, runtime):
    result = ctx.run(f"{runtime} image inspect {CRIMSON_FORGE_LOCAL_IMAGE}", warn=True, hide=True)
    if result.exited == 0:
        return CRIMSON_FORGE_LOCAL_IMAGE
    print(f"Local image {CRIMSON_FORGE_LOCAL_IMAGE} not found, pulling from ECR...")
    ctx.run(f"{runtime} pull {CRIMSON_FORGE_IMAGE}")
    return CRIMSON_FORGE_IMAGE

def _run_analysis(ctx, bin_file, fmt, out_file):
    workspace = Path.cwd()
    arch = bin_file.stem.rsplit(".", 1)[1]
    runtime = _container_runtime()
    image = _resolve_image(ctx, runtime)
    ctx.run(
        f"{runtime} run -ti -v {workspace}:/opt/workspace:Z --rm {image}"
        f" tools/analysis/graph.py -a {arch}"
        f" --log-level DEBUG"
        f" --skip-analysis"
        f" /opt/workspace/{bin_file}"
        f" {fmt}"
        f" /opt/workspace/{out_file}"
    )


@task
def clean(ctx):
    ctx.run(f"rm -rf {BIN_DIR}")
    ctx.run(f"rm -rf {GRAPHML_DIR}")
    ctx.run(f"rm -rf {GRAPHVIZ_DIR}")


@task
def asm(ctx):
    """Assemble all .asm files in src/ into bin/ using nasm, skipping up-to-date files."""
    BIN_DIR.mkdir(exist_ok=True)

    asm_files = list(SRC_DIR.glob("*.asm"))
    if not asm_files:
        print("No .asm files found in src/")
        return

    for asm_file in asm_files:
        out_file = BIN_DIR / (asm_file.stem + ".bin")

        if out_file.exists() and out_file.stat().st_mtime >= asm_file.stat().st_mtime:
            print(f"Skipping {asm_file} (up to date)")
            continue

        print(f"Assembling {asm_file} -> {out_file}")
        ctx.run(f"nasm -f bin {asm_file} -o {out_file}")


@task(pre=[asm])
def graphml(ctx):
    """Generate GraphML data for all .bin files in bin/, skipping up-to-date files."""
    GRAPHML_DIR.mkdir(exist_ok=True)

    bin_files = list(BIN_DIR.glob("*.bin"))
    if not bin_files:
        print("No .bin files found in bin/")
        return

    for bin_file in bin_files:
        parts = bin_file.stem.rsplit(".", 1)
        if len(parts) != 2:
            print(f"Skipping {bin_file} (filename must be <name>.<arch>.bin)")
            continue
        out_file = GRAPHML_DIR / (bin_file.stem + ".graphml")

        if out_file.exists() and out_file.stat().st_mtime >= bin_file.stat().st_mtime:
            print(f"Skipping {bin_file} (up to date)")
            continue

        print(f"Generating GraphML {bin_file} -> {out_file}")
        _run_analysis(ctx, bin_file, "graphml", out_file)


@task(pre=[asm])
def graphviz(ctx):
    """Generate Graphviz data for all .bin files in bin/, skipping up-to-date files."""
    GRAPHVIZ_DIR.mkdir(exist_ok=True)

    bin_files = list(BIN_DIR.glob("*.bin"))
    if not bin_files:
        print("No .bin files found in bin/")
        return

    for bin_file in bin_files:
        parts = bin_file.stem.rsplit(".", 1)
        if len(parts) != 2:
            print(f"Skipping {bin_file} (filename must be <name>.<arch>.bin)")
            continue
        out_file = GRAPHVIZ_DIR / (bin_file.stem + ".dot")

        if out_file.exists() and out_file.stat().st_mtime >= bin_file.stat().st_mtime:
            print(f"Skipping {bin_file} (up to date)")
            continue

        print(f"Generating Graphviz {bin_file} -> {out_file}")
        _run_analysis(ctx, bin_file, "graphviz", out_file)

@task(pre=[asm, graphml, graphviz])
def all(ctx):
    """Generat data for all formats."""
    pass


build = Collection('build')
build.add_task(all)
build.add_task(asm)
build.add_task(graphml)
build.add_task(graphviz)

ns = Collection()
ns.add_task(clean)
ns.add_collection(build)