import os
import subprocess
import sqlite3
import json

from viper.common.abstracts import Module
from viper.common.objects import File
from viper.core.database import Database
from viper.core.session import __sessions__

def mapClusters(c):
    return {
        "id": c[0],
        "description": c[1],
        "generation_level": c[2],
        # "last_update": c[3],
        # "graph": c[4],
        # "samples": c[5],
        "min_funcs": c[6],
        "max_funcs": c[7],
        # "dot": c[8],
        "tags": c[9],
    }

class Cluster(Module):
    cmd = 'cluster'
    description = 'Analysis and cluster malware'

    def __init__(self):
        super(Cluster, self).__init__()
        self.parser.add_argument('-a', '--analysis', action = 'store_true', help = 'Only analysis file')
        self.parser.add_argument('-c', '--cluster', action = 'store_true', help = 'Only execute clustering')
        self.parser.add_argument('-e', '--export', action = 'store_true', help = 'Export json from database')

    def run(self):
        super(Cluster, self).run()
        if self.args.analysis:
            self.analysis()
        elif self.args.cluster:
            self.cluster()
        elif self.args.export:
            self.export()
        else:
            self.analysis()
            self.cluster()
            self.export()

    def analysis(self):
        if not __sessions__.is_set():
            return
        
        file = __sessions__.current.file

        command = 'python /opt/cosa-nostra/pyew_batch.py '+ file.path
        process = subprocess.Popen(command, shell = True, stdout = subprocess.PIPE, stderr = subprocess.PIPE, cwd='/opt/cosa-nostra')
        out = process.stdout.read().decode('utf-8')
        err = process.stderr.read().decode('utf-8')
        process.stdout.close()
        process.stderr.close()

        self.log('success', 'Analyzer exited with output:')
        self.out(out)
        self.out(err)

    def cluster(self):
        command = 'python /opt/cosa-nostra/cn_clusterer.py'
        process = subprocess.Popen(command, shell = True, stdout = subprocess.PIPE, stderr = subprocess.PIPE)
        out = process.stdout.read().decode('utf-8')
        err = process.stderr.read().decode('utf-8')
        process.stdout.close()
        process.stderr.close()

        self.log('success', 'Clusterer exited with output:')
        self.out(out)
        self.out(err)

    def out(self, out):
        lines = out.split('\n')
        length = len(lines)

        for i in range(length):
            line = lines[i]
            if not line and i is length - 1:
                break

            self.log('item', line)

    def export(self):
        try:
            connection = sqlite3.connect('/opt/cosa-nostra/db.sqlite')
            cursor = connection.cursor()
            cursor.execute('SELECT * FROM clusters')

            results = cursor.fetchall()
            clusters = list(map(mapClusters, results))
            connection.close()
            
            with open('/root/.viper/static/malware-family/clusters.json', 'w') as f:
                json.dump(clusters, f, indent=4)

            self.log('success', 'Exported clusters')
        except:
            self.log('error', 'Failed to export')

