First release 0.0.8

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 <contracts folder>
+
+        --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()