Pull info from shodans internetdb

internetdb.py

#!/usr/bin/python3

import socket
import requests
import json
import argparse
from rich.console import Console


def get_arguments():
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--domain", "-d", help="Enter a domain e.g. tesla.com")
    group.add_argument("--file", "-f", help="Enter the file containing domains")
    return parser.parse_args()


SHODAN_BASE_URL = "https://internetdb.shodan.io/"
ARGS = get_arguments()
CONSOLE = Console()


def prep_domains() -> list:
    if ARGS.file:
        with open(ARGS.file, "r") as handle:
            domains = handle.readlines()
        domains = list(map(lambda c: c.strip(), domains))
    else:
        domains = [ARGS.domain]
    return domains


def get_ip_addresses(domains: list) -> list:
    IPs = set()
    for domain in domains:
        try:
            ip = socket.gethostbyname(domain)
            IPs.add(ip)
        except Exception as e:
            print(f"{domain} couldn't be resolved with error: {e}")
            continue
    return list(IPs)


def get_shodan_info(ip_list: list) -> list:
    info_list = []
    for ip in ip_list:
        full_url = SHODAN_BASE_URL + ip
        res = requests.get(full_url).json()
        if "No information available" in res:
            print(f"No information for {ip}")
            continue
        info_list.append(res)
    return info_list


def write_output(results: list):
    for element in results:
        try:
            filename = f"{element['ip'].replace('.', '_')}.json"
            with open(filename, "w") as handle:
                json.dump(element, handle, ensure_ascii=False, indent=4)
        except Exception as e:
            print(f"Failed to write output with error: {e}")
            continue


def main():
    with CONSOLE.status("") as status:
        domains = prep_domains()
        CONSOLE.print("[green]Prepared input!")
        status.update("[bold yellow]Resolving to IP addresses...")
        ips = get_ip_addresses(domains)
        CONSOLE.print(f"[green]Collected {len(ips)} unique IP addresses!")
        status.update("[bold yellow]Starting to pull info from shodan...")
        results = get_shodan_info(ips)
        CONSOLE.print("[green]Collected info from shodan!")
        write_output(results)
        CONSOLE.print("[green]Written output to respective files!")


if __name__ == '__main__':
    main()