vt_api.py

from flask import Flask, Blueprint, abort, request, url_for, jsonify
import requests
import virustotal
from os.path import join
from virustotal import VirusTotal
import config

vt = Blueprint('vt', __name__, url_prefix='/vt')

#### VIRUSTOTAL API ####

def getVTRoutes(vt):
    res = []
    endpoints = vt.getEndpoints()
    for base in endpoints:
        for idx in endpoints[base]['endpoints']:
            res.append(idx)
    return res

@vt.route('/', methods=["GET"])
def vt_main():
    vt = config.VT['VT']
    return jsonify({'endpoints': vt.getEndpoints()})

@vt.route('/file')
@vt.route('/file/')
def vt_file():
    vt = config.VT['VT']
    return jsonify({'endpoints': vt.File.getEndpoints()})

@vt.route('/file/report', methods=["GET"])
def vt_file_report():
    vt = config.VT['VT']
    args = request.args

    resource = args.get('resource')

    try:
        allinfo = bool(args.get('allinfo', False))
    except:
        return "Invalid value for allinfo, expected True/False!"

    try:
        res = vt.File.report(resource, allinfo=allinfo)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/scan', methods=["GET"])
def vt_file_scan():
    vt = config.VT['VT']
    args = request.args

    filename = args.get('file')
    if filename is not None:
        filename = join(config.UPLOAD_FOLDER, filename)

    specialurl = args.get('specialurl')

    try:
        res = vt.File.scan(filename, specialurl=specialurl)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/scan/uploadurl', methods=["GET"])
@vt.route('/file/scan/upload_url', methods=["GET"])
def vt_file_scan_uploadurl():
    vt = config.VT['VT']

    try:
        res = vt.File.scan_uploadurl()
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/rescan', methods=["GET"])
def vt_file_rescan():
    vt = config.VT['VT']
    args = request.args

    resource = args.get('resource')

    try:
        res = vt.File.rescan(resource)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/download', methods=["GET"])
def vt_file_download():
    vt = config.VT['VT']
    args = request.args

    hash = args.get('hash')

    try:
        res = vt.File.download(hash)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/behavior', methods=["GET"])
@vt.route('/file/behaviour', methods=["GET"])
def vt_file_behaviour():
    vt = config.VT['VT']
    args = request.args

    hash = args.get('hash')

    try:
        res = vt.File.behaviour(hash)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/networktraffic', methods=["GET"])
@vt.route('/file/network-traffic', methods=["GET"])
def vt_file_networktraffic():
    vt = config.VT['VT']
    args = request.args

    hash = args.get('hash')

    try:
        res = vt.File.networktraffic(hash)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/feed', methods=["GET"])
def vt_file_feed():
    vt = config.VT['VT']
    args = request.args

    timestamp = args.get('timestamp')

    try:
        res = vt.File.feed(timestamp)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/clusters', methods=["GET"])
def vt_file_clusters():
    vt = config.VT['VT']
    args = request.args

    date = args.get('date')

    try:
        res = vt.File.clusters(date)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/file/search', methods=["GET"])
def vt_file_search():
    vt = config.VT['VT']
    args = request.args

    query = args.get('query')

    offset = args.get('offset')

    try:
        res = vt.File.search(query, offset=offset)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/url')
@vt.route('/url/')
def vt_url():
    vt = config.VT['VT']
    return jsonify({'endpoints': vt.URL.getEndpoints()})

@vt.route('/url/report', methods=["GET"])
def vt_url_report():
    vt = config.VT['VT']
    args = request.args

    resource = args.get('resource')

    try:
        allinfo = bool(args.get('allinfo', False))
    except:
        return "Invalid value for allinfo, expected True/False!"
    
    try:
        scan = int(args.get('scan', 0))
    except:
        return "Invalid value for scan, expected int!"

    try:
        res = vt.URL.report(resource, allinfo=allinfo, scan=scan)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/url/scan', methods=["GET"])
def vt_url_scan():
    vt = config.VT['VT']
    args = request.args

    url = args.get('url') 

    try:
        res = vt.URL.scan(url)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/url/feed', methods=["GET"])
def vt_url_feed():
    vt = config.VT['VT']
    args = request.args

    timestamp = args.get('timestamp')

    try:
        res = vt.URL.feed(timestamp)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/domain')
@vt.route('/domain/')
def vt_domain():
    vt = config.VT['VT']
    return jsonify({'endpoints': vt.Domain.getEndpoints()})

@vt.route('/domain/report', methods=["GET"])
def vt_domain_report():
    vt = config.VT['VT']
    args = request.args

    domain = args.get('domain')

    try:
        res = vt.Domain.report(domain)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/ipaddress')
@vt.route('/ipaddress/')
@vt.route('/ip-address')
@vt.route('/ip-address/')
def vt_ipaddress():
    vt = config.VT['VT']
    return jsonify({'endpoints': vt.IpAddress.getEndpoints()})

@vt.route('/ipaddress/report', methods=["GET"])
@vt.route('/ip-address/report', methods=["GET"])
def vt_ipaddress_report():
    vt = config.VT['VT']
    args = request.args

    ip = args.get('ip')

    try:
        res = vt.IpAddress.report(ip)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/comments')
@vt.route('/comments/')
def vt_comments():
    vt = config.VT['VT']
    return jsonify({'endpoints': vt.Comments.getEndpoints()})

@vt.route('/comments/get', methods=["GET"])
def vt_comments_get():
    vt = config.VT['VT']
    args = request.args

    resource = args.get('resource')
    before = args.get('before')

    try:
        res = vt.Comments.get(resource, before=before)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

@vt.route('/comments/put', methods=["POST"])
def vt_comments_put():
    vt = config.VT['VT']
    args = request.args

    resource = args.get('resource')

    # comment = args.get('comment')
    # modify to accept data from post contents
    comment = None

    try:
        res = vt.Comments.put(resource, comment)
    except virustotal.VT_Error as msg:
        return jsonify({'error': str(msg)}), 500
    return jsonify({'results': res})

########################

def searchVT(query, offset=None):
    vt = config.VT['VT']

    try:
        res = vt.File.search(query, offset=offset)
    except virustotal.VT_Error as msg:
        print(msg)
        return {'error': str(msg)}
    return {"data": res}