<filename>malwarebazaar/commands/query.py<gh_stars>10-100
import click
from json import dumps
from rich.status import Status

from malwarebazaar.api import Bazaar
from malwarebazaar.output import *
from malwarebazaar.sample import Sample

ALLOWED_TYPES = ["hash", "imphash", "signature", "yara", "filetype", "clamav", "tag", "issuer", "subject", "tlsh",
                 "telfhash"]


@click.command(help="Queries the API for various file attributes.")
@click.option(
    "-l",
    "--limit",
    type=int,
    default=100,
    help="Limit the amount of samples returned (Default is 100).",
)
@click.option("-j", "--json", is_flag=True, help="Write JSON response to stdout.")
@click.option("-c", "--csv", is_flag=True, help="Convert output to CSV.")
@click.argument("query_type", type=click.Choice(ALLOWED_TYPES, case_sensitive=False))
@click.argument("query", type=str)
def query(limit, json, csv, query_type, query):
    c, ec = Console(), Console(stderr=True, style="bold red")
    bazaar = Bazaar(Config.get_instance()["api_key"])

    with Status("Querying MalwareBazaar..."):
        if query_type.lower() == "hash":
            res = bazaar.query_hash(query)
        elif query_type.lower() == "imphash":
            res = bazaar.query_imphash(query, limit=limit)
        elif query_type.lower() == "signature":
            res = bazaar.query_signature(query, limit=limit)
        elif query_type.lower() == "yara":
            res = bazaar.query_yara(query, limit=limit)
        elif query_type.lower() == "filetype":
            res = bazaar.query_filetype(query, limit=limit)
        elif query_type.lower() == "clamav":
            res = bazaar.query_clamav_signature(query, limit=limit)
        elif query_type.lower() == "tag":
            res = bazaar.query_tag(query, limit=limit)
        elif query_type.lower() == "issuer":
            res = bazaar.query_signing_issuer(query)
        elif query_type.lower() == "subject":
            res = bazaar.query_signing_subject(query)
        elif query_type.lower() == "tlsh":
            res = bazaar.query_tlsh(query, limit=limit)
        elif query_type.lower() == "telfhash":
            res = bazaar.query_telfhash(query, limit=limit)

    if res["query_status"] != "ok":
        ec.print(f"Invalid Bazaar response: {res['query_status']}")
        exit(-1)

    if json:
        c.print(dumps(res["data"], indent=2))
    elif csv:
        samples = [Sample(**sample) for sample in res["data"]]
        c.print("md5,sha1,sha256,imphash,signature")
        for s in samples:
            c.print(
                f"{s.md5_hash},{s.sha1_hash},{s.sha256_hash},{s.imphash},{s.signature}"
            )
    else:
        samples = [Sample(**sample) for sample in res["data"]]
        multiple_samples(samples, c)

<filename>malwarebazaar/commands/query.py<gh_stars>10-100
import click
from json import dumps
from rich.status import Status

from malwarebazaar.api import Bazaar
from malwarebazaar.output import *
from malwarebazaar.sample import Sample

ALLOWED_TYPES = ["hash", "imphash", "signature", "yara", "filetype", "clamav", "tag", "issuer", "subject", "tlsh",
                 "telfhash"]


@click.command(help="Queries the API for various file attributes.")
@click.option(
    "-l",
    "--limit",
    type=int,
    default=100,
    help="Limit the amount of samples returned (Default is 100).",
)
@click.option("-j", "--json", is_flag=True, help="Write JSON response to stdout.")
@click.option("-c", "--csv", is_flag=True, help="Convert output to CSV.")
@click.argument("query_type", type=click.Choice(ALLOWED_TYPES, case_sensitive=False))
@click.argument("query", type=str)
def query(limit, json, csv, query_type, query):
    c, ec = Console(), Console(stderr=True, style="bold red")
    bazaar = Bazaar(Config.get_instance()["api_key"])

    with Status("Querying MalwareBazaar..."):
        if query_type.lower() == "hash":
            res = bazaar.query_hash(query)
        elif query_type.lower() == "imphash":
            res = bazaar.query_imphash(query, limit=limit)
        elif query_type.lower() == "signature":
            res = bazaar.query_signature(query, limit=limit)
        elif query_type.lower() == "yara":
            res = bazaar.query_yara(query, limit=limit)
        elif query_type.lower() == "filetype":
            res = bazaar.query_filetype(query, limit=limit)
        elif query_type.lower() == "clamav":
            res = bazaar.query_clamav_signature(query, limit=limit)
        elif query_type.lower() == "tag":
            res = bazaar.query_tag(query, limit=limit)
        elif query_type.lower() == "issuer":
            res = bazaar.query_signing_issuer(query)
        elif query_type.lower() == "subject":
            res = bazaar.query_signing_subject(query)
        elif query_type.lower() == "tlsh":
            res = bazaar.query_tlsh(query, limit=limit)
        elif query_type.lower() == "telfhash":
            res = bazaar.query_telfhash(query, limit=limit)

    if res["query_status"] != "ok":
        ec.print(f"Invalid Bazaar response: {res['query_status']}")
        exit(-1)

    if json:
        c.print(dumps(res["data"], indent=2))
    elif csv:
        samples = [Sample(**sample) for sample in res["data"]]
        c.print("md5,sha1,sha256,imphash,signature")
        for s in samples:
            c.print(
                f"{s.md5_hash},{s.sha1_hash},{s.sha256_hash},{s.imphash},{s.signature}"
            )
    else:
        samples = [Sample(**sample) for sample in res["data"]]
        multiple_samples(samples, c)
