chapinb / chickadee-rs

Yet another GeoIP enrichment tool
1 stars 1 forks source link

Implement a Resolver for VirusTotal #6

Open onsali opened 1 year ago

onsali commented 1 year ago

Hi Chapin, I wrote this resolver for the VirusTotal API which I am looking to integrate into Chickadee

use reqwest::Client;
use serde_json::Value;

//hardcoded api key; implement alternative
const VIRUS_TOTAL_API_KEY: &str = "YOUR API KEY HERE";

async fn lookup_ip(ip: &str) -> Result<(), reqwest::Error> {
    let client = Client::new();

    let url = format!("https://www.virustotal.com/api/v3/ip_addresses/{}", ip);

    //create request headers from API key
    let mut headers = HeaderMap::new();
    headers.insert(
        "x-apikey",
        HeaderValue::from_str(VIRUS_TOTAL_API_KEY).expect("Invalid API key"),
    );

    //send async request
    let mut response = client
        .get(&url)
        .headers(headers)
        .send()
        .await?;

    //parse response body as json
    let response_body: Value = response.json().await?;

    println!("Response: {:?}", response_body);

    Ok(())
}

#[tokio::main]
async fn main() {
    //implement cli arg here
    let ip_address = "204.77.151.204";

    if let Err(err) = lookup_ip(ip_address).await {
        eprintln!("Error: {}", err);
    }
}

/*to do:
parse output
impl arg for choosing resolver
*/
chapinb commented 1 year ago

Thanks for sharing! Here's some tips for next steps on integration.

Unfortunately chickadee-rs doesn't yet support async, so we will need to remove that component for now, though I hope to bring it into the framework soon.

The lookup_ip() function can go in a new module, such as "resolver/virustotal.rs".

You will want to update the function to return a Vec<String>, where the String is serialized JSON, so that it matches the return value of the IP API (in the future we could look at building a custom serializer struct, but can skip that for now)

Next you will want to update the resolve_ip_address() function to include this new virustotal resolver, maybe with a new argument influence a conditional to pick between the two available resolvers.

I believe the last step is to add a new command line argument that influences the conditional to select the correct resolver. Alternatively, we could look up from both sources all the time.

chapinb commented 1 year ago

A few other notes to be aware of:

  1. The script does allow the user to filter on columns. For the VT resolver see if it is easy enough to leverage the existing code by specifying a few fields in the VT output and seeing if you can filter the output.
  2. Does this API endpoint allow enrichment of multiple IPs at once? If so, you may need to update the lookup_ip() function to allow multiple IPs. If not, we'll want to create a loop to generate a Vec<> of results.
  3. The API key could be stored in an environment variable for now, which will make it easier to read the contents into the script. In the future it can be read from a configuration file.
  4. Try to avoid .expect() where you can and handle/return errors. I am using the anyhow::Result() and you could use the anyhow::anyhow macro to raise errors.
onsali commented 1 year ago

I have rewritten the program without asynch, and managed to integrate it using an extra CLI argument for defining the resolver type. I have opened a pull request #7 under the branch virustotal_implementation, please take a look.