Skip to content

Commit

Permalink
chore(examples): add python example
Browse files Browse the repository at this point in the history
  • Loading branch information
Jon-Becker committed Sep 2, 2024
1 parent 6649fb2 commit 262baf3
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 0 deletions.
27 changes: 27 additions & 0 deletions examples/python/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
# Example: Invoking Heimdall via Python

This Python script demonstrates how to use the `heimdall` CLI tool to decompile smart contracts via Python. It provides a simple class structure to define arguments and manage the decompilation process, with support for customizing various decompilation options.

## Overview

The script utilizes the `heimdall decompile` command to decompile a target contract and retrieve its Solidity and Yul source code, along with its ABI. For ease of use, the script abstracts the command-line interface of `heimdall` into a Python class, allowing users to easily invoke the decompilation process programmatically.

### Key Components

- **DecompileArgs**: A class to define the arguments for the decompilation process, such as the target contract address, RPC URL, and other flags.
- **DecompiledContract**: A class to represent the output of the decompilation, including the Solidity/Yul source code and ABI.
- **Decompiler**: A class that abstracts the `heimdall decompile` command, handling argument parsing and command execution.
- **is_heimdall_installed**: A utility function to check if the `heimdall` CLI tool is installed on the system.

## Usage

1. **Install `heimdall`**: Ensure that the `heimdall` CLI tool is installed on your system.

```bash
which heimdall

2. **Run the Script**: Execute the Python script to decompile a target contract.

```bash
python main.py
```
126 changes: 126 additions & 0 deletions examples/python/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
import subprocess
from typing import Optional
import tempfile
import os
import json

"""
arguments to be passed to `heimdall decompile`
"""
class DecompileArgs:
def __init__(
self,
target,
rpc_url="",
default=False,
skip_resolving=False,
include_solidity=False,
include_yul=False,
timeout=10000,
):
self.target = target
self.rpc_url = rpc_url
self.default = default
self.skip_resolving = skip_resolving
self.include_solidity = include_solidity
self.include_yul = include_yul
self.timeout = timeout


"""
decompiled contract
"""
class DecompiledContract:
def __init__(self, path):
# walk the directory and print the files
for root, dirs, files in os.walk(path):
for file in files:
full_path = os.path.join(root, file)

# read contents of the file
# if .sol or .yul, read the contents as a string
# if .json, read the contents as a dict
with open(full_path, "r") as f:
contents = f.read()

if file.endswith(".sol") or file.endswith(".yul"):
self.source = contents
elif file.endswith(".json"):
self.abi = json.loads(contents)


"""
abstraction on top of `heimdall decompile`
"""
class Decompiler:
def __init__(self, args: DecompileArgs):
self.args = args

def decompile(self) -> Optional[DecompiledContract]:
try:
# generate a temp dir
temp_dir = tempfile.TemporaryDirectory()

command = [
"heimdall",
"decompile",
self.args.target,
"--output",
temp_dir.name,
]

if self.args.rpc_url:
command.extend(["--rpc-url", self.args.rpc_url])
if self.args.default:
command.append("--default")
if self.args.skip_resolving:
command.append("--skip-resolving")
if self.args.include_solidity:
command.append("--include-sol")
if self.args.include_yul:
command.append("--include-yul")
if self.args.timeout:
command.extend(["--timeout", str(self.args.timeout)])

subprocess.check_output(command)
return DecompiledContract(temp_dir.name)
except Exception as e:
print("Error: ", e)
return None


"""
checks if the `heimdall` cli is installed on the system via `which heimdall`.
"""
def is_heimdall_installed():
try:
subprocess.check_output(["which", "heimdall"])
return True
except subprocess.CalledProcessError:
return False


def main():
if not is_heimdall_installed():
print("heimdall does not seem to be installed on your system.")
print("please install heimdall before running this script.")
return

args = DecompileArgs(
target="0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2",
rpc_url="https://eth.llamarpc.com",
include_solidity=True,
)

decompiled_contract = Decompiler(args).decompile()

if decompiled_contract:
print("Decompiled Contract:")
print("Source:")
print(decompiled_contract.source)
print("ABI:")
print(decompiled_contract.abi)


if __name__ == "__main__":
main()

0 comments on commit 262baf3

Please sign in to comment.