import.py

import os
import sqlite3
import json
import shutil
import hashlib
from virustotal_python import Virustotal

vtotal = Virustotal("Insert API key here")
cwd = "/home/pixelweaver/.cuckoo"

init_db = False
if not os.path.isfile('samples.db'):
    init_db = True
else:
    if os.path.isfile('samples.bak'):
        os.remove('samples.bak')
    shutil.copyfile('samples.db', 'samples.bak')

conn = sqlite3.connect('samples.db')
c = conn.cursor()

if init_db:
    c.execute('''CREATE TABLE samples
             (rowid INTEGER PRIMARY KEY, sha256 TEXT, mark_count INTEGER)''')
    c.execute('''CREATE TABLE snapshots
                 (sample_id INTEGER, 
                 cpu_user REAL, cpu_system REAL, cpu_idle REAL, cpu_interrupt REAL, cpu_dpc REAL, 
                 v_mem_total INTEGER, v_mem_available INTEGER, v_mem_used INTEGER, v_mem_free INTEGER,
                 s_mem_total INTEGER, s_mem_used INTEGER, s_mem_free INTEGER, s_mem_per REAL, s_mem_sin INTEGER, s_mem_sout INTEGER,
                 d_io_read_count INTEGER, d_io_write_count INTEGER, d_io_read_bytes INTEGER, d_io_write_bytes REAL, d_io_read_time INTEGER, d_io_write_time INTEGER,
                 n_io_bytes_sent INTEGER, m_io_bytes_recv INTEGER, n_io_packets_sent INTEGER, n_io_packets_recv INTEGER,
                 highest_pid INTEGER, proc_number INTEGER)''')
    conn.commit()

try:
    analyses_dir = [entry
                    for entry in next(os.walk(os.path.join((os.path.join(cwd, "storage/analyses")))))[1]
                    if entry.isnumeric()
                    and os.path.isfile(
                        os.path.join((os.path.join(cwd, "storage/analyses", entry, "files", "metrics.json"))))
                    and os.path.isfile(
                        os.path.join((os.path.join(cwd, "storage/analyses", entry, "reports", "report.json"))))]

    unknown_from_virustotal = []
    for analysis_dir in analyses_dir:
        report_fp = open(os.path.join(cwd, "storage/analyses", analysis_dir, "reports", "report.json"))
        report_json = json.load(report_fp)
        sample_id = None

        if "virustotal" in report_json and "sha256" in report_json["virustotal"]:  # If VirusTotal report fetching was correctly conducted in cuckoo, use it
            c.execute('INSERT INTO samples (sha256, mark_count) VALUES (?, ?)',
                      (report_json["virustotal"]["sha256"], report_json["virustotal"]["positives"]))
            sample_id = c.lastrowid
        else:  # If not, fetch report now
            sha256_hash = hashlib.sha256()
            with open(os.path.join(cwd, "storage/analyses", analysis_dir, "binary"), "rb") as f:
                # Read and update hash string value in blocks of 4K
                for byte_block in iter(lambda: f.read(4096), b""):
                    sha256_hash.update(byte_block)
            resp = vtotal.file_report([sha256_hash.hexdigest()])["json_resp"]
            if resp["response_code"] is 0:
                unknown_from_virustotal.append(analysis_dir)
                c.execute('INSERT INTO samples (sha256, mark_count) VALUES (?, ?)',
                          (sha256_hash.hexdigest(), -1))
            else:
                c.execute('INSERT INTO samples (sha256, mark_count) VALUES (?, ?)',
                          (resp["sha256"], resp["positives"]))
            sample_id = c.lastrowid

        snapshots_fp = open(os.path.join(cwd, "storage/analyses", analysis_dir, "files", "metrics.json"))
        snapshots_json = json.load(snapshots_fp)

        for snapshot_json in snapshots_json:
            c.execute(
                'INSERT INTO snapshots VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)',
                (sample_id,
                 snapshot_json[0][0], snapshot_json[0][1], snapshot_json[0][2], snapshot_json[0][3],
                 snapshot_json[0][4],
                 snapshot_json[1][0], snapshot_json[1][1], snapshot_json[1][3], snapshot_json[1][4],
                 snapshot_json[2][0], snapshot_json[2][1], snapshot_json[2][2], snapshot_json[2][3],
                 snapshot_json[2][4], snapshot_json[2][5],
                 snapshot_json[3][0], snapshot_json[3][1], snapshot_json[3][2], snapshot_json[3][3],
                 snapshot_json[3][4], snapshot_json[3][5],
                 snapshot_json[4][0], snapshot_json[4][1], snapshot_json[4][2], snapshot_json[4][3],
                 max(snapshot_json[4]), len(snapshot_json[4])
                 ))

    conn.commit()
    conn.close()
    print("Added " + str(len(analyses_dir)) + " exploitable results found in CWD to the sample database.")
    if os.path.isfile('samples.bak'):
        os.remove('samples.bak')

    if len(unknown_from_virustotal) != 0:
        print("...Alas, " + str(len(unknown_from_virustotal)) + " files were unknown from virus total, they are :")
        print(unknown_from_virustotal)
        print("They have been marked by -1 mark_count in the database.")
except Exception as e:
    conn.close()
    os.remove('samples.db')
    if os.path.isfile('samples.bak'):
        os.rename('samples.bak', 'samples.db')

    print("An error occurred while ingesting results : " + str(e))
    print("Changes to the database were reverted.")