Simple security lookups via CLI
#!/usr/bin/env python3

import click
from seclook.lookups import (
    shodan_lookup,
    virustotal_lookup,
    emailrep_lookup,
    abuseipdb_lookup,
    greynoise_lookup,
    threatfox_lookup,
    pulsedive_lookup,
)
from seclook.openai import gpt4_summarize
import json
import os


@click.command()
@click.argument("service", required=False)
@click.argument("value", required=False)
@click.option("--export", is_flag=True, help="Export JSON to Desktop")
@click.option("--gpt4", is_flag=True, help="Use GPT-4 to summarize results")
def main(service, value, export, gpt4):
    """Perform lookups from various security services

    - Use `seclook [service] [value]` to perform a lookup.

    - Use `seclook list` to see a list of available services.
    """

    services = [
        "list",  # Keep at top, so services.pop(0) always removes it
        "shodan",
        "virustotal",
        "emailrep",
        "abuseipdb",
        "greynoise",
        "threatfox",
        "pulsedive",
    ]

    if not service:
        raise click.UsageError("Missing service argument.")

    if service.lower() not in services:
        raise click.UsageError(f"'{service}' is not available in seclook.")

    if service.lower() == "list":
        services.pop(0)  # L29
        click.echo("Available services:")
        for service in services:
            click.echo("- " + service.capitalize())
        click.echo("Run 'seclook [service] [value]' to perform a lookup.")
        return

    if not value:
        raise click.UsageError(f"Missing value argument for '{service}'.")

    if service.lower() == "shodan":
        result = shodan_lookup.search(value)
    elif service.lower() == "virustotal":
        result = virustotal_lookup.search(value)
    elif service.lower() == "emailrep":
        result = emailrep_lookup.search(value)
    elif service.lower() == "abuseipdb":
        result = abuseipdb_lookup.search(value)
    elif service.lower() == "greynoise":
        result = greynoise_lookup.search(value)
    elif service.lower() == "threatfox":
        result = threatfox_lookup.search(value)
    elif service.lower() == "pulsedive":
        result = pulsedive_lookup.search(value)

    if export:
        desktop = os.path.join(os.path.expanduser("~"), "Desktop")
        filename = os.path.join(desktop, f"seclook_{service}_{value}.json")

        with open(filename, "w") as outfile:
            json.dump(result, outfile, indent=4)

        click.echo(f"Results exported to {filename}")
        return

    if gpt4:
        result = gpt4_summarize.search(service, result)
        pretty_result = json.dumps(result, indent=4)
        click.echo(pretty_result)

    pretty_result = json.dumps(result, indent=4)
    click.echo(pretty_result)


if __name__ == "__main__":
    main()

seclook/cli.py
ackatz/seclook
import configparser
from seclook.find_config import find_config
import textwrap
import sys


def load_config(service_name, key_required=True):
    try:
        config_path = find_config()
        config = configparser.ConfigParser()
        config.read(config_path)
    except configparser.ParsingError:
        print(
            textwrap.dedent(
                f"""
        ParsingError: Ensure the following lines are present in config.ini:

        [{service_name}]
        api_key = <Note: A key is{' required' if key_required
                else ' optional'} to use this API>
        """
            )
        )
        sys.exit()

    if key_required:
        try:
            api_key = config[service_name]["api_key"]
        except KeyError:
            print(
                textwrap.dedent(
                    f"""
            KeyError: Ensure the following lines are present in config.ini:

            [{service_name}]
            api_key = <Note: A key is required to use this API>
            """
                )
            )
            sys.exit()
    else:
        try:
            api_key = config[service_name].get("api_key", "")
        except KeyError:
            print(
                textwrap.dedent(
                    f"""
            KeyError: Ensure the following lines are present in config.ini:

            [{service_name}]
            api_key = <Note: A key is optional to use this API>
            """
                )
            )
            sys.exit()

    return config, api_key

seclook/config_helper.py
ackatz/seclook
import os


def find_config():
    dirs_to_check = [
        os.getcwd(),
        os.path.dirname(os.path.abspath(__file__)),
        os.path.join(os.path.expanduser("~"), ".seclook"),
    ]

    for directory in dirs_to_check:
        config_path = os.path.join(directory, "config.ini")
        if os.path.exists(config_path):
            return config_path

    raise FileNotFoundError(
        "No config.ini file found in any of the checked directories:\n"
        "({})".format(", ".join(dirs_to_check)) + "\n"
        "Copy and rename config.ini.sample as config.ini into "
        "any of the directories mentioned above:\n"
        "https://github.com/ackatz/seclook/blob/main/config.ini.sample"
    )

seclook/find_config.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://api.abuseipdb.com/api/v2/check"


def search(value):
    config, abuseipdb_api_key = load_config("abuseipdb")
    headers = {"Accept": "application/json", "Key": abuseipdb_api_key}
    params = {"ipAddress": value, "maxAgeInDays": 90, "verbose": True}
    response = requests.get(base_url, headers=headers, params=params)
    return response.json()

seclook/lookups/abuseipdb_lookup.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://emailrep.io/{}"


def search(value):
    config, emailrep_api_key = load_config("emailrep")
    headers = {"Key": emailrep_api_key}
    response = requests.get(base_url.format(value), headers=headers)
    return response.json()

seclook/lookups/emailrep_lookup.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://api.greynoise.io/v3/community/{}"


def search(value):
    config, greynoise_api_key = load_config("greynoise", key_required=False)
    headers = {"Accept": "application/json", "key": greynoise_api_key}
    response = requests.get(base_url.format(value), headers=headers)
    return response.json()

seclook/lookups/greynoise_lookup.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://pulsedive.com/api/info.php"


def search(value):
    config, pulsedive_api_key = load_config("pulsedive", key_required=False)
    headers = {"Accept": "application/json"}
    response = requests.get(
        base_url.format(value),
        headers=headers,
        params={"key": pulsedive_api_key, "indicator": value},
    )
    return response.json()

seclook/lookups/pulsedive_lookup.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://api.shodan.io/shodan/host/{}"


def search(value):
    config, shodan_api_key = load_config("shodan", key_required=False)
    headers = {"Accept": "application/json"}
    response = requests.get(
        base_url.format(value), headers=headers, params={"key": shodan_api_key}
    )
    return response.json()

seclook/lookups/shodan_lookup.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://threatfox-api.abuse.ch/api/v1/"


def search(value):
    config, threatfox_api_key = load_config("threatfox", key_required=False)
    headers = {
        "Content-Type": "application/json",
        "API-KEY": threatfox_api_key,
    }
    payload = {"query": "search_ioc", "search_term": value}
    response = requests.post(base_url, headers=headers, json=payload)
    return response.json()

seclook/lookups/threatfox_lookup.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://www.virustotal.com/api/v3/search"


def search(value):
    config, virustotal_api_key = load_config("virustotal")
    headers = {"x-apikey": virustotal_api_key}
    params = {"query": value}
    response = requests.get(base_url, headers=headers, params=params)
    return response.json()

seclook/lookups/virustotal_lookup.py
ackatz/seclook
import requests
from seclook.config_helper import load_config

base_url = "https://api.openai.com/v1/chat/completions"


def search(service, json_response):
    config, openai_api_key = load_config("openai")
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + openai_api_key,
    }
    payload = {
        "model": "gpt-4",
        "messages": [
            {
                "role": "user",
                "content": f"You are a security analyst summarizer. "
                f"Given the following JSON response from {service}"
                f", summarize the response to tell me what I need to know: "
                f"{json_response}",
            }
        ],
    }
    response = requests.post(base_url, headers=headers, json=payload)
    return response.json()

seclook/openai/gpt4_summarize.py
ackatz/seclook
from click.testing import CliRunner
from seclook.cli import main


def test_export():
    runner = CliRunner()
    result = runner.invoke(main, ["threatfox", "1.1.1.1", "--export"])
    assert "Results exported to" in result.output

seclook/tests/test_export.py
ackatz/seclook
from click.testing import CliRunner
from seclook.cli import main


def test_list():
    runner = CliRunner()
    result = runner.invoke(main, ["list"])
    assert result.exit_code == 0
    assert "Available services:" in result.output

seclook/tests/test_list.py
ackatz/seclook
from click.testing import CliRunner
from seclook.cli import main


def test_missing_service_and_value():
    runner = CliRunner()
    result = runner.invoke(main, [])
    assert result.exit_code != 0
    assert "Error: Missing service argument." in result.output

seclook/tests/test_missing_service_value.py
ackatz/seclook
from click.testing import CliRunner
from seclook.cli import main


def test_threatfox_missing_value():
    runner = CliRunner()
    result = runner.invoke(main, ["threatfox"])
    assert result.exit_code != 0
    assert "Missing value argument for 'threatfox'." in result.output


def test_pulsedive_missing_value():
    runner = CliRunner()
    result = runner.invoke(main, ["pulsedive"])
    assert result.exit_code != 0
    assert "Missing value argument for 'pulsedive'." in result.output


def test_shodan_missing_value():
    runner = CliRunner()
    result = runner.invoke(main, ["shodan"])
    assert result.exit_code != 0
    assert "Missing value argument for 'shodan'." in result.output


def test_virustotal_missing_value():
    runner = CliRunner()
    result = runner.invoke(main, ["virustotal"])
    assert result.exit_code != 0
    assert "Missing value argument for 'virustotal'." in result.output


def test_emailrep_missing_value():
    runner = CliRunner()
    result = runner.invoke(main, ["emailrep"])
    assert result.exit_code != 0
    assert "Missing value argument for 'emailrep'." in result.output


def test_abuseipdb_missing_value():
    runner = CliRunner()
    result = runner.invoke(main, ["abuseipdb"])
    assert result.exit_code != 0
    assert "Missing value argument for 'abuseipdb'." in result.output


def test_greynoise_missing_value():
    runner = CliRunner()
    result = runner.invoke(main, ["greynoise"])
    assert result.exit_code != 0
    assert "Missing value argument for 'greynoise'." in result.output

seclook/tests/test_missing_value.py
ackatz/seclook
from click.testing import CliRunner
from seclook.cli import main


def test_unknown_service():
    runner = CliRunner()
    result = runner.invoke(main, ["unknown", "value"])
    assert result.exit_code != 0
    assert "is not available in seclook" in result.output

seclook/tests/test_unknown_service.py
ackatz/seclook
from click.testing import CliRunner
from seclook.cli import main


def test_threatfox_valid_value():
    runner = CliRunner()
    result = runner.invoke(main, ["threatfox", "1.1.1.1"])
    assert result.exit_code == 0


def test_pulsedive_valid_value():
    runner = CliRunner()
    result = runner.invoke(main, ["pulsedive", "1.1.1.1"])
    assert result.exit_code == 0


def test_shodan_valid_value():
    runner = CliRunner()
    result = runner.invoke(main, ["shodan", "1.1.1.1"])
    assert result.exit_code == 0


def test_virustotal_valid_value():
    runner = CliRunner()
    result = runner.invoke(main, ["virustotal", "1.1.1.1"])
    assert result.exit_code == 0


def test_emailrep_valid_value():
    runner = CliRunner()
    result = runner.invoke(main, ["emailrep", "andrew@akatz.org"])
    assert result.exit_code == 0


def test_abuseipdb_valid_value():
    runner = CliRunner()
    result = runner.invoke(main, ["abuseipdb", "1.1.1.1"])
    assert result.exit_code == 0


def test_greynoise_valid_value():
    runner = CliRunner()
    result = runner.invoke(main, ["greynoise", "1.1.1.1"])
    assert result.exit_code == 0

seclook/tests/test_valid_service_value.py
ackatz/seclook
from setuptools import setup, find_packages

setup(
    name="seclook",
    author="Andrew Katz",
    author_email="andrew@akatz.org",
    license="MIT",
    version="0.5.1",
    packages=find_packages(),
    include_package_data=True,
    install_requires=[
        "click",
        "requests",
        "configparser",
    ],
    entry_points={
        "console_scripts": [
            "seclook=seclook.cli:main",
        ],
    },
    description="Simple security lookups via CLI",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    url="https://github.com/ackatz/seclook",
)

setup.py
ackatz/seclook