Immunefi Program Data Generator

genscope.py

#!/usr/bin/env python3

from typing import List
from sys import argv, exit, stderr
from enum import Enum
from argparse import ArgumentParser
import json

try:
    import requests
except ImportError as e:
    print("[!] Package 'requests' not installed or found. Please install: pip install requests")
    exit(1)

from xml.etree.ElementTree import ElementTree, Element

try:
    from defusedxml.ElementTree import parse, fromstring
except ImportError as e:
    print("[!] Package 'defusedxml' not installed or found. Please install: pip install defusedxml")
    exit(1)


IMMUNEFI_BASE_URL = "https://immunefi.com"
IMMUNEFI_PROJECT_PATH = "/bounty"
VERBOSE = 0
DEBUG = True

def eprint(*args, **kwargs):
    print(*args, file=stderr, **kwargs)

def debug(*args, **kwargs):
    if VERBOSE > 1:
        verbose(*args, prefix="[**] ", **kwargs)

def verbose(msg, prefix="[*] ", *args, **kwargs):
    if VERBOSE > 0:
        msg = prefix + msg
        eprint(*args, **kwargs)

class PayoutEnum(str, Enum):
    BLOCKCHAIN = "BLOCKCHAIN"
    WEB = "WEB"
    OTHER = "OTHER"

#TODO: cleanup
BLOCKCHAIN_PAYOUT_TITLES = [
    "Smart Contracts and Blockchain" 
]
BLOCKCHAIN_PAYOUT_TITLES = [s.lower() for s in BLOCKCHAIN_PAYOUT_TITLES]

def payout_enum(text: str) -> PayoutEnum:
    if text.lower() in BLOCKCHAIN_PAYOUT_TITLES:
        return PayoutEnum.BLOCKCHAIN
    return PayoutEnum.OTHER

class Payout:
    def __init__(self, level, payout: str, type_: PayoutEnum, sort: int):
        self.level: str = level
        self.payout: str = payout
        self.type: PayoutEnum = type_
        self.sort: int = sort

    def to_dict(self):
        return {"payout": self.payout, "level": self.level, "type": self.type, "sort": self.sort}

    def __repr__(self):
        return f"<Payout(level={self.level},payout={self.payout},type={self.type},sort={self.sort})>"

class AssetEnum(str, Enum):
    SMART_CONTRACT = "SMART_CONTRACT"
    WEB = "WEB"
    DAPP = "DAPP"
    OTHER = "OTHER"


class Asset:
    def __init__(self, type_: AssetEnum, target: str, name: str):
        self.type = type_
        self.target = target
        self.name = name
    
    def to_dict(self):
        return {"target": self.target, "type": self.type, "name": self.name}

    def __repr__(self):
        return f"<Asset(type={self.type},target={self.target},name={self.name})>"

class Program:
    def __init__(self, payouts: List[Payout], assets: List[Asset]):
        self.payouts = payouts
        self.assets = assets
    def to_json(self):
        payouts = []
        assets = []
        for payout in self.payouts:
            payouts.append(payout.to_dict())
        for asset in self.assets:
            assets.append(asset.to_dict())
        return json.dumps({"payouts": payouts, "assets": assets})

def build_program_url(slug: str):
    return IMMUNEFI_BASE_URL + IMMUNEFI_PROJECT_PATH + "/" + slug


def ele_has_child(ele: Element, data: str, tag: str):
    data = data.replace(" ", "").replace("\n", "").lower()
    for child in ele:
        if child.tag.lower() == tag.lower():
            val = child.text
            if val:
                val = val.replace(" ", ""). replace("\n", "").lower()
            if val and data in val:
                return True
    return False

def parse_payouts_from_section(section: Element):
    payouts = []
    # sections contain h3, p and divs
    for subsection in section:
        current_type = PayoutEnum.OTHER
        # find type which is a p sibling to the main div holding table
        if subsection.tag.lower() == "p":
            if len(subsection) == 1:
                if subsection[0].tag.lower() == "strong" and subsection[0].text:
                    current_type = payout_enum(subsection[0].text.lower())
        # match on next div holding table
        # may be more than one table for each type, but should be only divs
        if subsection.tag.lower() == "div":
            # browser parses diffrently, there'a  div wrapping abunch fo meta
            # stuff that needs to be avoided
            if len(subsection) > 0:
                if subsection[0].tag.lower() != "dl":
                    continue
            # dl holds two divs
            #  first is div with one dd (for level)
            #  second dt is just string "level", skip it
            #  second div holds payout amount
            #    get amt from first dd
            index = 0
            for payout_dl in subsection:
                if payout_dl.tag == "dl":
                    level = "Unknown"
                    amount = "Unknown"
                    if len(payout_dl) != 2:
                        eprint("Err on payout_dl length", len(payout_dl))
                        continue
                    else:
                        level_div = payout_dl[0]
                        if len(level_div) != 2:
                            eprint("Err on leve_div length")
                            continue
                        level_dd = level_div[0]
                        amt_div = payout_dl[1]
                        if len(amt_div) != 2:
                            eprint("Err on amt div")
                            continue
                        amt_dd = amt_div[0]

                        level = level_dd.text or amount
                        amount = amt_dd.text or amount
                        payouts.append(Payout(level, amount, current_type,index))
                        index += 1
    return payouts


def parse_assets_from_section(section: Element):
    assets = []
    if len(section) > 2:
        for asset_dl in section[2]:
            type_ = AssetEnum.OTHER
            name = "Generic"
            target = ""
            if len(asset_dl) != 2:
                eprint("Err on asset dl")
                continue
            target_div = asset_dl[0]
            if len(target_div) != 2:
                eprint("Err on asset target div")
                continue
            target_dd = target_div[0]
            if len(target_dd) != 1:
                eprint("Err on asset target dd")
                continue
            target = target_dd[0].text or ""

            type_div = asset_dl[1]
            if len(type_div) != 2:
                eprint("Err on asset type div")
                continue
            type_dd = target_div[0]
            type_str = type_dd.text or ""
            if type_str.lower().startswith("smart contract -"):
                type_ = AssetEnum.SMART_CONTRACT
                name = type_str.lower().split("-", 1)[1].strip()
            asset = Asset(type_, target, name)
            assets.append(asset)
    return assets




SECTIONS_XPATH = ".//body/div/div/main/section/article/div/section"
def build_scope(slug: str):
    verbose("Downloading data from Immunefi")
    res = requests.get(build_program_url(slug))
    if res.status_code != 200:
        return 1
    tree: Element = fromstring(res.text)
    payouts = []
    assets = []
    for section in tree.iterfind(SECTIONS_XPATH):
        if ele_has_child(section, "Rewards by Threat Level", "h3"):
            payouts = parse_payouts_from_section(section)

        if ele_has_child(section, "Assets In Scope", "h3"):
            assets = parse_assets_from_section(section)

    program = Program(payouts, assets)
    data = program.to_json()
    print(data)

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-s", "--slug", required=True)
    args = parser.parse_args()
    exit(build_scope(args.slug))