#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import getpass
import argparse
import readline
import json


import rlcompleter
import atexit
import glob

# tab completion
def complete(text, state):
    return (glob.glob(text+'*')+[None])[state]
readline.set_completer_delims(' \t\n;')
readline.parse_and_bind("tab: complete")
readline.set_completer(complete)

# history file
histfile = os.path.join(os.environ['HOME'], '.vxcage_history')
try:
    readline.read_history_file(histfile)
except IOError:
    pass
atexit.register(readline.write_history_file, histfile)
del histfile, readline, rlcompleter




try:
    import requests
    from progressbar import *
    from prettytable import PrettyTable
except ImportError as e:
    sys.exit("ERROR: Missing dependency: %s" % e)



def color(text, color_code):
    return '\x1b[%dm%s\x1b[0m' % (color_code, text)

def cyan(text):
    return color(text, 36)

def bold(text):
    return color(text, 1)

def logo():
    print("")
    print(cyan("  `o   O o   O .oOo  .oOoO' .oOoO .oOo. "))
    print(cyan("   O   o  OoO  O     O   o  o   O OooO' "))
    print(cyan("   o  O   o o  o     o   O  O   o O     "))
    print(cyan("   `o'   O   O `OoO' `OoO'o `OoOo `OoO' "))
    print(cyan("                                O       "))
    print(cyan("                             OoO' ") + " by nex")
    print("")

def help():
    print("Available commands:")
    print("  " + bold("tags") + "         Retrieve list of tags")
    print("  " + bold("find") + "         Query a file by md5, sha256, ssdeep, imphash, tag or date")
    print("  " + bold("get") + "          Download a file by sha256")
    print("  " + bold("dump") + "         Dump a list of md5, sha256, ssdeep hashes")
    print("  " + bold("add") + "          Upload a file to the server")
    print("  " + bold("last") + "         Retrieve a list of the last x files uploaded")
    print("  " + bold("total") + "        Total number of samples")
    print("  " + bold("stats") + "        File type stats")
    print("  " )
    print("  " + bold("version") + "      Version of remote vxcage server")
    print("  " + bold("license") + "      Print the software license")
    print("  " )
    print("  " + bold("help | ?") + "         Show this help")
    print("  " + bold("exit | quit") + "  Exit cli application")


class VxCage(object):
    def __init__(self, host, port, xmock, ssl=False, auth=False):
        self.host = host
        self.port = port
        self.ssl = ssl
        self.auth = auth
        self.xmock = xmock
        self.username = None
        self.password = None

    def authenticate(self):
        if self.auth:
            self.username = raw_input("Username: ")
            self.password = getpass.getpass("Password: ")

    def build_url(self, route):
        if self.ssl:
            url = "https://"
            if self.port is None:
                self.port = 443
        else:
            if self.port is None:
                self.port = 8080
            url = "http://"

        url += "%s:%s%s%s" % (self.host, self.port, self.xmock, route)

        return url

    def check_errors(self, code):
        if code == 400:
            print("ERROR: Invalid request format")
            return True
        elif code == 500:
            print("ERROR: Unexpected error, check your server logs")
            return True
        else:
            return False

    def tags_list(self):
        req = requests.get(self.build_url("/tags/list"),
                           auth=(self.username, self.password),
                           verify=False)
        try:
            res = req.json()
        except:
            try:
                res = req.json
            except Exception as e:
                print("ERROR: Unable to parse results: {0}".format(e))
                return

        if self.check_errors(req.status_code):
            return

        table = PrettyTable(["tag"])
        table.align = "l"
        table.padding_width = 1

        for tag in res:
            table.add_row([tag])

        print(table)
        print("Total: %s" % len(res))

    def dump_list(self, hType):
        req = requests.get(self.build_url("/malware/dump/"+hType),
                           auth=(self.username, self.password),
                           verify=False)
        try:
            res = req.json()
        except:
            try:
                res = req.json
            except Exception as e:
                print("ERROR: Unable to parse results: {0}".format(e))
                return

        if self.check_errors(req.status_code):
            return

        table = PrettyTable([hType])
        table.align = "l"
        table.padding_width = 1

        for hType in res:
            table.add_row(hType)

        print(table)
        print("Total: %s" % len(res))

    def malware_total(self):