generator.py

import argparse
import helper.libc as libc
from functools import partial


def partial_format(fmt, **kwargs):
    return partial(fmt.format, kwargs)


base_call = """
#include <time.h>
#include <stddef.h>
#include <stdlib.h>
{headers}

{ra} {name}({args})
{{{{
    {init}
    {{crash}}
    return {call};
}}}}
"""

libcFunctions = {
    "malloc": libc.LibCFunction(
        "malloc", "size", "size_t", "void *", headers=["stdlib.h"]
    ),
    "fopen": libc.LibCFunction(
        "fopen",
        "path, mode",
        "const char *, const char *",
        "FILE *",
        headers=["stdio.h"],
    ),
    "realloc": libc.LibCFunction(
        "realloc", "ptr, size", "void *, size_t", "void *", headers=["stdlib.h"]
    ),
    "calloc": libc.LibCFunction(
        "calloc", "nmemb, size", "size_t, size_t", "void *", headers=["stdlib.h"]
    ),
    "strtok": libc.LibCFunction(
        "strtok", "str, delim", "char *, const char *", "char *", headers=["string.h"]
    ),
    "close": libc.LibCFunction("close", "fd", "int", "int", headers=["unistd.h"]),
    "fclose": libc.LibCFunction("fclose", "fd", "FILE *", "int", headers=["stdio.h"]),
    "read": libc.LibCFunction(
        "read", "fd, buf, count", "int, void *, size_t", "ssize_t", headers=["unistd.h"]
    ),
    "write": libc.LibCFunction(
        "write",
        "fd, buf, count",
        "int, const void *, size_t",
        "ssize_t",
        headers=["unistd.h"],
    ),
    "fread": libc.LibCFunction(
        "fread",
        "ptr, size, nmemb, stream",
        "void *, size_t, size_t, FILE *",
        "size_t",
        headers=["stdio.h"],
    ),
    "fwrite": libc.LibCFunction(
        "fwrite",
        "ptr, size, nmemb, stream",
        "const void *, size_t, size_t, FILE *",
        "size_t",
        headers=["stdio.h"],
    ),
}


bases = {
    "malloc": base_call.format(
        ra="void *",
        name="malloc",
        args="size_t size",
        init=libcFunctions["malloc"].load(),
        call=libcFunctions["malloc"].call(),
        headers=libcFunctions["malloc"].header(),
    ),
    "fopen": base_call.format(
        ra="FILE *",
        name="fopen",
        args="const char *path, const char *mode",
        init=libcFunctions["fopen"].load(),
        call=libcFunctions["fopen"].call(),
        headers=libcFunctions["fopen"].header(),
    ),
    "realloc": base_call.format(
        ra="void *",
        name="realloc",
        args="void *ptr, size_t size",
        init=libcFunctions["realloc"].load(),
        call=libcFunctions["realloc"].call(),
        headers=libcFunctions["realloc"].header(),
    ),
    "calloc": base_call.format(
        ra="void *",
        name="calloc",
        args="size_t nmemb, size_t size",
        init=libcFunctions["calloc"].load(),
        call=libcFunctions["calloc"].call(),
        headers=libcFunctions["calloc"].header(),
    ),
    "strtok": base_call.format(
        ra="char *",
        name="strtok",
        args="char *str, const char *delim",
        init=libcFunctions["strtok"].load(),
        call=libcFunctions["strtok"].call(),
        headers=libcFunctions["strtok"].header(),
    ),
    "close": base_call.format(
        ra="int",
        name="close",
        args="int fd",
        init=libcFunctions["close"].load(),
        call=libcFunctions["close"].call(),
        headers=libcFunctions["close"].header(),
    ),
    "fclose": base_call.format(
        ra="int",
        name="fclose",
        args="FILE *fd",
        init=libcFunctions["fclose"].load(),
        call=libcFunctions["fclose"].call(),
        headers=libcFunctions["fclose"].header(),
    ),
    "read": base_call.format(
        ra="ssize_t",
        name="read",
        args="int fd, void *buf, size_t count",
        init=libcFunctions["read"].load(),
        call=libcFunctions["read"].call(),
        headers=libcFunctions["read"].header(),
    ),
    "write": base_call.format(
        ra="ssize_t",
        name="write",
        args="int fd, const void *buf, size_t count",
        init=libcFunctions["write"].load(),
        call=libcFunctions["write"].call(),
        headers=libcFunctions["write"].header(),
    ),
    "fread": base_call.format(
        ra="size_t",
        name="fread",
        args="void *ptr, size_t size, size_t nmemb, FILE *stream",
        init=libcFunctions["fread"].load(),
        call=libcFunctions["fread"].call(),
        headers=libcFunctions["fread"].header(),
    ),
    "fwrite": base_call.format(
        ra="size_t",
        name="fwrite",
        args="const void *ptr, size_t size, size_t nmemb, FILE *stream",
        init=libcFunctions["fwrite"].load(),
        call=libcFunctions["fwrite"].call(),
        headers=libcFunctions["fwrite"].header(),
    ),
}

errors = {
    "malloc": "return NULL;",
    "fopen": "return NULL;",
    "realloc": "return NULL;",
    "calloc": "return NULL;",
    "strtok": "return NULL;",
    "close": "return -1;",
    "fclose": "return -1;",
    "read": "return -1;",
    "write": "return -1;",
    "fread": "return 0;",
    "fwrite": "return 0;",
}

rand = """
    static int init = 0;
    if (!init)
    {{
        srand(time(NULL));
        init = 1;
    }}

    if (rand() % {rand_count} == 0)
    {{
        {error};
    }}
"""

after = """
    static int count = 0;
    if (count++ > {fail_after})
    {{
        {error};
    }}
"""


class Generator:
    def __init__(self, name: str, rand=None, always=False, fail_after=None):
        self.rand = rand
        self.base = bases[name]
        self.error = errors[name]
        self.always = always
        self.fail_after = fail_after

    def random(self):
        return rand.format(rand_count=self.rand, error=self.error)

    def after(self):
        return after.format(fail_after=self.fail_after, error=self.error)

    def generate(self):
        crash = ""
        if self.always:
            crash = self.error
        elif self.rand:
            crash = self.random()
        elif self.fail_after:
            crash = self.after()
        return self.base.format(crash=crash)


def main():
    parser = argparse.ArgumentParser(description="Generate a crashing function")

    parser.add_argument("--rand", type=int, help="Add random crash to function")
    parser.add_argument("--always", action="store_true", help="Always crash")
    parser.add_argument("--name", type=str, help="Name of the function")
    parser.add_argument("--all", action="store_true", help="Generate all functions")
    parser.add_argument("--fail-after", type=int, help="Fail after N calls")

    args = parser.parse_args()
    if args.name and not args.all:
        gen = Generator(
            args.name, rand=args.rand, always=args.always, fail_after=args.fail_after
        )
        res = gen.generate()
        print(res)
    elif args.all:
        for name in bases:
            gen = Generator(
                name, rand=args.rand, always=args.always, fail_after=args.fail_after
            )
            res = gen.generate()
            print(res)
        return


if __name__ == "__main__":
    main()