diff --git a/audit_helper.py b/audit_helper.py new file mode 100644 index 0000000..d14c7c1 --- /dev/null +++ b/audit_helper.py @@ -0,0 +1,311 @@ +#!/usr/bin/python3 + +# SPDX-License-Identifier: MIT + +import sys +import os + +ver = "0.0.8" +testPrefix = "" +testFolder = "test/" + +boilerPlate = '\nimport "forge-std/Test.sol";\n' + +# version() print version +def version(): + print(f"Solidity Contract Audit Helper v" + ver) + +# Search for contracts and foundry test contracts +def locateSolidityFiles(folder): + tests = [] + files = [] + for file in os.listdir(folder): + if file.endswith(".sol"): + files.append(file) + + for file in os.listdir(testFolder): + if file.endswith(".t.sol"): + tests.append(file) + + return files, tests + +def checkFoundry(): + foundry = os.system("forge --version") + if foundry == 0: + print(f"Foundry located.") + else: + print(f"Foundry not found!\nInstall via: https://github.com/foundry-rs/foundry#installation") + +def cleanupFoundry(): + print("Removing Foundry init contract files...") + files = ["script/Counter.s.sol","src/Counter.sol","test/Counter.t.sol"] + + for f in files: + try: + os.remove(f) + except OSError as e: + print("Error: %s - %s." % (e.filename, e.strerror)) + + +def npmInit(): + for x in os.listdir(): + if x == "package.json": + # Install packages if required with existing hardhat install + npm_packages = os.system("npm install") + if npm_packages != 0: + print("[!] ERROR: Problem installing extra packages") + exit() + else: + return + + print("[?] INFO: Package.json not found, not installing further npm packages") + +def setupFoundry(folder, testFolder, oz, sm): + # Check for foundry install + checkFoundry() + + # forge init + print("Initialising Foundry ...\n") + + forge_init = os.system("forge init --force --vscode --no-commit --no-git") + if forge_init != 0: + print("[!] ERROR: Problem initialising Foundry") + exit() + + # Install OpenZeppelin repo + if oz == True: + try: + os.system("forge install openzeppelin/openzeppelin-contracts --no-git") + except: + print("[!] ERROR: Problem installing OpenZeppelin-Contracts") + + # Install Solmate repo + if sm == True: + try: + os.system("forge install transmissions11/solmate --no-git") + except: + print("[!] ERROR: Problem installing Solmate") + + # Clean up Foundry init Counter files + cleanupFoundry() + + npmInit() + + # Rewrite the new foundry.toml for to include -c and -o values + output = [] + packages = os.path.isdir("node_modules") + toml = open("foundry.toml", "r") + + for line in toml: + if line.find("src =") > -1: + output.append("src = '" + folder + "'\n") + output.append("test = '" + testFolder + "'\n") + elif line.find("libs =") > -1 and packages == True: + # include node_modules + output.append("libs = ['lib', 'node_modules']\n") + else: + output.append(line) + toml.close() + + # Overwrite existing foundry.toml with modified data + toml = open("foundry.toml", "w") + toml.writelines(output) + toml.close() + + # Update remappings.txt + remappings = os.system("forge remappings > remappings.txt") + if remappings != 0: + print("[!] ERROR: Problem remapping packages") + exit() + else: + if oz == True: os.system("echo \"openzeppelin-contracts/=lib/openzeppelin-contracts/contracts/\" >> remappings.txt") + if sm == True: os.system("echo \"solmate/=lib/solmate/src//\" >> remappings.txt") + +# Pull all public and external functions from a contract +def worker(folder, inFile): + output = [] + closer = " \n\t}\n" + wait = False + + file = str(folder + inFile) + readFile = open(file, 'r') + + header = '''// SPDX-License-Identifier: UNLICENSED +/// @author @HardlyCodeMan https://github.com/HardlyCodeMan/audit_helper/ +/// @info Boilerplate test file auto generated by Solidity Contract Audit Helper v''' + ver + "\n" + + setUp = "\tfunction setUp() public {\n\n\t}\n" + + output.append(header) + + lines = readFile.readlines() + for line in lines: + line = line.strip() + + # Skip lines that start with // || /* || * || * + if line.startswith("//") or line.startswith("/*") or line.startswith("*") or line.startswith(" *"): + continue + # Locate lines that contain pragma || function && public || external + if line.find("pragma") > -1: + output.append(line + "\n") + output.append(boilerPlate) + elif line.find("import") > -1: + if line.find("./") > -1: + line = line.replace("./", "../" + folder) + output.append(line + "\n") + elif line.find("contract") > -1 or wait == True: + # If inheritance is over multiple lines + if wait == True: + if line.strip().endswith("{"): + output.append(line.strip() + "\n") + wait = False + else: + output.append("\t" + line.strip() + "\n") + else: + output.append("\n") + # Append contract name with "Test" and include "is Test" + # Split by default splits at the space character + contractLine = line.split() + contractLine[1] = str(contractLine[1]) + "Test" + + # Is does the contract inherit from any other contracts + if contractLine[2] == "is": + contractLine[2] = "is Test," + else: + contractLine[2] = "is Test {" + + modifiedLine = "" + for i in range(len(contractLine)): + modifiedLine = modifiedLine + contractLine[i] + " " + + output.append(modifiedLine + "\n") + + if modifiedLine.strip().endswith("{") == False: + wait = True + + # Add the setUp() function for foundry tests + output.append(setUp) + + elif (line.find("function") > -1 and (line.find("public") > -1 or line.find("external") > -1)): + contractLine = line.split() + funcName = str(contractLine[1]) + funcName = funcName.capitalize() + contractLine[1] = "test" + str(funcName) + + modifiedLine = "\t" + for i in range(len(contractLine)): + modifiedLine = modifiedLine + contractLine[i] + " " + + output.append(modifiedLine + "\n") + + if line.endswith("{"): + output.append(closer) + + # Add the contract closing } + output.append("\n}") + + readFile.close() + + if len(output) > 1: + writeTests(output, inFile) + else: + print("[?] INFO: No tests to write.") + +# Write to a new test file, overwrite if existing so utilise a prefix where needed +def writeTests(input, file): + filename = file.split(".") + outFile = str(testFolder + testPrefix + os.path.splitext(file)[0] + ".t.sol") + + print(f"[<-] Writing test: " + testPrefix + filename[0] + ".t.sol") + write = open(outFile, 'w') + write.writelines(input) + write.close() + +# worker() main control loop +def main(folder, testFolder): + # Append trailing / to folders if required + if folder.find("/") == -1: + folder = folder + "/" + if testFolder.find("/") == -1: + testFolder = testFolder + "/" + + print(f"\nWorking dir: " + str(folder)) + files, tests = locateSolidityFiles(folder) + + # Print contract files + if len(files) >= 1: + print(f"\nLocated contracts: ") + for file in range(len(files)): + print(f" [c] " + files[file]) + else: + print("\nNo contracts found.") + + print("\nWorking dir: " + str(testFolder)) + # Print test files + if len(tests) >= 1: + print(f"\nLocated tests: ") + for file in range(len(tests)): + print(f" [t] " + tests[file]) + else: + print("\nNo test found.") + + if len(files) == 0: + print("\nNothing to do.") + exit() + + # Start working + for file in range(len(files)): + print("\n[->] Working on contract: " + files[file]) + worker(folder, files[file]) + +if __name__ == "__main__": + # Sort cli arguments + args = sys.argv + runSetup = False + run = False + oz = False + sm = False + + # No flags sent + if len(args) == 1: + version() + print(""" + Usage: + audit_helper -c + + --help | -h : Display this menu + --version | -v : Print the version number + --contracts | -c : Location of contracts + --output | -o : Test output folder. (Default: test/) + --setup | -s : Initialise Foundry project and edit foundry.toml accordingly + --openzeppelin | -oz : Requires -s. Initialize with OpenZeppelin-Contracts repo + --solmate | -sm : Requires -s. Initialize with Solate repo + --prefix | -p : Test file prefix (Default: "") + """) + else: + version() + for i in range(len(args)): + if args[i] == "--version" or args[i] == "-v": + version() + exit() + if args[i] == "--setup" or args[i] == "-s": + runSetup = True + if args[i] == "--openzeppelin" or args[i] == "-oz": + oz = True + if args[i] == "--solmate" or args[i] == "-sm": + sm = True + if args[i] == "--output" or args[i] == "-o": + testFolder = args[i +1] + if args[i] == "--prefix" or args[i] == "-p": + testPrefix = args[i +1] + if args[i] == "--contracts" or args[i] == "-c": + contracts = args[i +1] + run = True + + if run == True and runSetup == True: + setupFoundry(contracts, testFolder, oz, sm) + main(contracts, testFolder) + elif run == True: + main(contracts, testFolder) + + exit() \ No newline at end of file