check

#!/opt/pwn.college/python

import pwnlib.context
import pwnlib.asm
import pwnlib.elf
import contextlib
import capstone
import tempfile
import socket
import magic
import time
import sys
import os

pwnlib.context.context.arch = "amd64"
cs = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)

sys.path.append(os.path.dirname(__file__)+"/.py")
import chal #pylint:disable=import-error,wrong-import-position

if os.geteuid() == 0:
    os.seteuid(65534)

class ChallengeFailed(Exception):
    pass
class ChallengeFailedPrint(ChallengeFailed):
    pass

def print_prompt():
    print(f"""hacker@{socket.gethostname()}:{
        os.getcwd().replace(os.path.expanduser('~'), '~', 1)
    }$ """, end="", flush=True)

def slow_print(what):
    for c in what:
        print(c, end="", flush=True)
        if "FAST" not in os.environ:
            time.sleep(0.05)
    print("")

def dramatic_command(command, actual_command=None):
    print_prompt()
    slow_print(command)
    exit_code = os.WEXITSTATUS(
        os.system(command if actual_command is None else actual_command)
    )
    if "FAST" not in os.environ:
        time.sleep(0.5)
    return exit_code

@contextlib.contextmanager
def disable_fd(fd):
    bkup = os.dup(fd)
    os.close(fd)
    yield
    os.dup2(bkup, fd)
    os.close(bkup)

def _assemble(asm):
    with disable_fd(2):
        return pwnlib.asm.asm(asm)

def assemble(asm):
    if not getattr(chal, "allow_asm", False):
        raise ChallengeFailedPrint(
            "This challenge requires you to assemble the code yourself. "
            "Please do that."
        )

    try:
        x86 = _assemble(asm)
    except pwnlib.exception.PwnlibException as e:
        errors = e.message.split("\n")
        if (
            errors[0].startswith("There was an error running") and
            "Assembler messages" in errors[3]
        ):
            msg = "\n".join(
                "- "+line.split(" ", 1)[-1]
                for line in errors[4:]
                if line
            )
            raise ChallengeFailedPrint(
                f"Your assembly did not assemble cleanly. The errors:\n{msg}"
            ) from e

        raise ChallengeFailedPrint(
            "Your assembly resulted in unexpected assembly errors:\n" +
            e.message
        ) from e

    return x86

def get_raw_binary(content):
    m = magic.from_buffer(content)
    if m == "ELF 64-bit relocatable":
        raise ChallengeFailedPrint(
            "You provided an ELF object file, rather than an actual\n"
            "ELF executable. The object file is an intermediate result\n"
            "of the compilation/assembly process. Please 'link' it into an\n"
            "executable by using:\n"
            "\n"
            "hacker@dojo:~$ ld final.elf your-object-file.o"
            "\n"
            "This will create a 'final.elf' file that you can pass to this program!"
        )
    if (
        "ELF 64-bit LSB shared object" in m or
        "ELF 64-bit LSB pie executable" in m or
        "ELF 64-bit LSB executable" in m
    ):
        if not getattr(chal, "allow_elf", True):
            raise ChallengeFailedPrint(
                "This challenge requires you to provide raw binary code, "
                "but you provided an ELF. Please extract your .text segment "
                "using 'objcopy'!"
            )
        tmpelf = tempfile.mktemp()
        with open(tmpelf, "wb") as o:
            o.write(content)
        text = pwnlib.elf.ELF(tmpelf).get_section_by_name(".text")
        if not text:
            raise ChallengeFailedPrint(
                "The ELF you provided is missing the .text section!"
            )
        rawbin = text.data()
        if not rawbin:
            raise ChallengeFailedPrint(
                "The .text section of the ELF you provided is empty!"
            )
        return rawbin
    elif "text" in m:
        return assemble(content.decode('latin1')+"\n")
    else: #assuming this is binary code
        return content

@contextlib.contextmanager
def output_color(color, stream=sys.stdout):
    CODES = {
        'HEADER': '\033[95m',
        'OKBLUE': '\033[94m',
        'OKCYAN': '\033[96m',
        'OKGREEN': '\033[92m',
        'WARNING': '\033[93m',
        'FAIL': '\033[91m',
        'ENDC': '\033[0m',
        'BOLD': '\033[1m',
        'UNDERLINE': '\033[4m',
    }
    print(CODES[color.upper()], end="", flush=True, file=stream)
    try:
        yield
    finally:
        print(CODES['ENDC'], end="", flush=True, file=stream)

def do_check(stage_id, stage_description, *args, **kwargs):
    try:
        if not hasattr(chal, stage_id):
            return None

        print("")
        with output_color("header"):
            print(getattr(chal, stage_id+"_prologue", f"{stage_description}...").format(**chal.__dict__))
        r = getattr(chal, stage_id)(*args, **kwargs)
        with output_color("okgreen"):
            print(getattr(chal, stage_id+"_success", "... YES! Great job!").format(**chal.__dict__))
        return r
    except (ChallengeFailed, AssertionError) as _e:
        with output_color("fail"):
            print(getattr(chal, stage_id+"_failure", "... oops, we found an issue! Details below:\n").format(**chal.__dict__))
            print(_e)
        raise

def main():
    if len(sys.argv) == 2:
        filename = sys.argv[1]
        try:
            #pylint:disable=consider-using-with,unspecified-encoding
            content = open(filename, "rb").read()
        except FileNotFoundError:
            print(f"File {filename} not found.")
            return
        except PermissionError:
            print(f"Permission denied when opening {filename}.")
            return
    elif not os.isatty(0):
        content = sys.stdin.buffer.read()
    else:
        if chal.allow_asm:
            print("Please input your assembly. Press Ctrl+D when done!")
        else:
            print(f"Please run this program as `{sys.argv[0]} your-program.elf`!")
            sys.exit(1)
        content = b""
        while line := sys.stdin.buffer.readline():
            content += line

    raw_binary = get_raw_binary(content)

    disas = list(cs.disasm(raw_binary, 0))
    if not disas:
        print("Your binary failed to disassemble.")
        return

    num_instructions = getattr(chal, "num_instructions", None)
    if num_instructions is not None and len(disas) != num_instructions:
        print(
            f"This challenge expects {num_instructions} "
            f"instruction{'s' if num_instructions != 1 else ''}, "
            f"but you provided {len(disas)}."
        )
        return

    if hasattr(chal, "assembly_prefix"):
        binary_prefix = _assemble(chal.assembly_prefix)
        raw_binary = binary_prefix + raw_binary
    if hasattr(chal, "assembly_suffix"):
        binary_suffix = _assemble(chal.assembly_suffix)
        raw_binary = raw_binary + binary_suffix

    do_check("check_raw_binary", "Checking the binary code", raw_binary)
    do_check("check_disassembly", "Checking the assembly code", disas)

    elf_filename = getattr(chal, 'final_filename', "/tmp/your-program")
    os.rename(
        pwnlib.asm.make_elf(raw_binary, extract=False),
        elf_filename
    )

    if os.geteuid() == 65534:
        os.seteuid(0)

    do_check("check_runtime", "Checking runtime behavior", elf_filename)

    if getattr(chal, "give_flag", False):
        print("")
        print("Here is your flag!")
        #pylint:disable=consider-using-with,unspecified-encoding
        print(open("/flag").read())


if __name__ == '__main__':
    try:
        main()
    except (ChallengeFailedPrint) as _e:
        print("Check failed:")
        print()
        with output_color("FAIL"):
            print(_e)
    except (ChallengeFailed, AssertionError):
        # we assume we already printed
        pass
    except Exception as _e: #pylint:disable=broad-exception-caught
        print("Unexpected error during challenge evaluation! The error:")
        print(_e)
        raise