fenrir.py

import os
from multiprocessing import Pool, cpu_count

import click
from mythril.mythril import SolidityContract, Report, SymExecWrapper, fire_lasers

import resource


@click.command()
@click.option('--contract-directory', '-i', default="./input")
@click.option('--output-directory', '-o', default="./output")
@click.option('--cores', default=4)
@click.option('--memory', default=8)
def main(contract_directory, output_directory, cores, memory):
    rsrc = resource.RLIMIT_DATA
    soft, hard = resource.getrlimit(rsrc)
    resource.setrlimit(rsrc, (memory * 1073741824, hard))

    fenrir(contract_directory, output_directory, cores)


def analyze_contract(input_output):
    path, output = input_output
    contract = SolidityContract(path)

    issues = fire_lasers(SymExecWrapper(contract, "", "dfs"), None)

    report = Report()
    for issue in issues:
        report.append_issue(issue)

    if len(issues) > 0:
        with open(output, 'w+') as output_file:
            output_file.write(report.as_markdown())
            output_file.flush()


def fenrir(contract_directory: str, output_directory: str, cores: int):
    if cores > cpu_count():
        print("[*] Your computer is not BEAST enough for that")
        return

    pool = Pool(cores)

    if not os.path.isdir(contract_directory):
        print("[*] Contract directory doesn't exist")
        return

    if not os.path.isdir(output_directory):
        print("[*] Output directory doesn't exist, creating ...")
        os.mkdir(output_directory)

    def names(filename):
        return "{}/{}".format(contract_directory, filename), "{}/{}.md".format(output_directory, filename)

    contract_files = list(map(names, os.listdir(contract_directory)))

    pool.map(analyze_contract, contract_files)

    pool.close()
    pool.join()


if __name__ == "__main__":
    main()