diff --git a/conf/maldb.db b/conf/maldb.db index 8b0acc9..af67de3 100644 Binary files a/conf/maldb.db and b/conf/maldb.db differ diff --git a/imports/db_handler.py b/imports/db_handler.py index b36d6fa..4e210d6 100644 --- a/imports/db_handler.py +++ b/imports/db_handler.py @@ -20,8 +20,10 @@ class DBHandler: return self.cur.execute("SELECT ID, TYPE, LANGUAGE, ARCHITECTURE, PLATFORM, NAME FROM Malwares").fetchall() def get_mal_names(self): + # Sqlite3 returns a tuple even if a single value is returned # We use x[0] for x to unpack the tuples + return [val[0] for val in self.cur.execute("SELECT NAME FROM Malwares").fetchall()] def query(self, query, param=''): @@ -35,3 +37,20 @@ class DBHandler: except lite.Error as e: print "An error occurred:", e.args[0] sys.exit() + + def close_connection(self): + try: + self.cur.close() + self.con.close() + return + except lite.Error as e: + print "An error occurred:", e.args[0] + sys.exit() + + def renew_connection(self): + try: + self.con = lite.connect(globals.vars.db_path) + self.cur = self.con.cursor() + except lite.Error as e: + print "An error occurred:", e.args[0] + sys.exit() diff --git a/imports/globals.py b/imports/globals.py index 212facf..c8a8169 100644 --- a/imports/globals.py +++ b/imports/globals.py @@ -42,7 +42,6 @@ class init: eula_file = conf_folder + '/eula_run.conf' maldb_ver_file = conf_folder + '/db.ver' giturl = 'https://github.com/ytisf/theZoo/blob/master' - addrs = ['reverce_tcp/', 'crazy_mal/', 'mal/', 'show malwares'] class bcolors: @@ -69,7 +68,7 @@ class vars: ############ DEBUGGING ############### #### SET TO ZERO BEFORE COMMIT ####### - + # DEBUG_LEVEL 0 = NO DEBUGGING # DEBUG_LEVEL 1 = DEBUG DOWNLOADS # DEBUG_LEVEL 2 = DEBUG SQL QUERIES @@ -106,14 +105,14 @@ class vars: eula_file = conf_folder + '/eula_run.conf' maldb_ver_file = conf_folder + '/db.ver' db_path = conf_folder + "/maldb.db" - giturl_dl = 'https://github.com/ytisf/theZoo/raw/master/malwares/' + giturl_dl = 'https://github.com/ytisf/theZoo/raw/master/' giturl = 'https://github.com/ytisf/theZoo' with file(maldb_ver_file) as f: db_ver = f.read() # ASCII Art is a must... - screen = random.randrange(1, 3) + screen = random.randrange(1, 4) if screen is 1: maldb_banner = "\n" @@ -152,5 +151,12 @@ class vars: maldb_banner += " `-. | __.-'\n" maldb_banner += " `-|__.--'\n" - addrs = ['reverce_tcp/', 'crazy_mal/', 'mal/', 'show malwares'] - addrs = ['list', 'search', 'get', 'exit'] + elif screen is 3: + maldb_banner = " __ ___ __ ____ ____\n" + maldb_banner += " / |/ /___ _/ / ______ _________ / __ \/ __ )\n" + maldb_banner += " / /|_/ / __ `/ / | /| / / __ `/ ___/ _ \______/ / / / __ |\n" + maldb_banner += " / / / / /_/ / /| |/ |/ / /_/ / / / __/_____/ /_/ / /_/ /\n" + maldb_banner += " /_/ /_/\__,_/_/ |__/|__/\__,_/_/ \___/ /_____/_____/\n\n" + maldb_banner += " version: " + version + "\n" + maldb_banner += " db_version: " + db_ver + "\n" + maldb_banner += " built by: " + authors + "\n\n" diff --git a/imports/terminal_handler.py b/imports/terminal_handler.py index 2638b25..7105319 100644 --- a/imports/terminal_handler.py +++ b/imports/terminal_handler.py @@ -1,10 +1,9 @@ -import csv import sys import re import globals from imports import manysearches -from imports.updatehandler import Updater +from imports.update_handler import Updater from imports import db_handler @@ -61,7 +60,6 @@ class Controller: # Checks if normal or freestyle search if re.match('^search', cmd): manySearch = manysearches.MuchSearch() - num_args = len(cmd.rsplit(' ')) try: args = cmd.rsplit(' ')[1:] manySearch.sort(args) @@ -73,8 +71,9 @@ class Controller: sys.exit(1) if cmd == 'update-db': - updateHandler = Updater() - updateHandler.get_maldb_ver() + update_handler = Updater() + db_ver = update_handler.get_maldb_ver() + update_handler.update_db(db_ver) self.MainMenu() if cmd == 'report-mal': @@ -116,11 +115,11 @@ class Controller: self.MainMenu() if cmd == 'get': - updateHandler = Updater() - #try: - updateHandler.get_malware(self.currentmodule) - #except: - print globals.bcolors.RED + '[-] ' + globals.bcolors.WHITE + 'Error getting malware.' + update_handler = Updater() + try: + update_handler.get_malware(self.currentmodule) + except: + print globals.bcolors.RED + '[-] ' + globals.bcolors.WHITE + 'Error getting malware.' self.MainMenu() # If used the 'use' command if re.match('^use', cmd): diff --git a/imports/updatehandler.py b/imports/update_handler.py similarity index 61% rename from imports/updatehandler.py rename to imports/update_handler.py index 86c9956..d1102af 100644 --- a/imports/updatehandler.py +++ b/imports/update_handler.py @@ -17,6 +17,7 @@ # along with this program. If not, see . import sys +from os import remove, rename import urllib2 from imports import globals from imports import db_handler @@ -24,6 +25,9 @@ from imports import db_handler class Updater: + def __init__(self): + self.db = db_handler.DBHandler() + def get_maldb_ver(self): ''' Get current malwareDB version and see if we need an update @@ -36,69 +40,62 @@ class Updater: "No malware DB version file found.\nPlease try to git clone the repository again.\n") return 0 - def update_db(self): + def update_db(self, curr_db_version): ''' Just update the database from GitHub :return: ''' - try: - with file(globals.vars.maldb_ver_file) as f: - f = f.read() - except IOError: - print( - "No malware DB version file found.\nPlease try to git clone the repository again.\n") - return 0 - - curr_maldb_ver = f + if globals.vars.DEBUG_LEVEL is 1: + print locals() response = urllib2.urlopen( - globals.vars.giturl_dl_dl + globals.vars.maldb_ver_file) + globals.vars.giturl_dl + globals.vars.maldb_ver_file) new_maldb_ver = response.read() - if new_maldb_ver == curr_maldb_ver: - print globals.bcolors.GREEN + '[+]' + globals.bcolors.WHITE + " No need for an update.\n" + globals.bcolors.GREEN + '[+]' + globals.bcolors.WHITE + " You are at " + new_maldb_ver + " which is the latest version." - sys.exit(1) + if new_maldb_ver == curr_db_version: + print globals.bcolors.GREEN + '[+]' + globals.bcolors.WHITE + " theZoo is up to date :)\n" + globals.bcolors.GREEN + '[+]' + globals.bcolors.WHITE + " You are at " + new_maldb_ver + " which is the latest version." + return + + print globals.bcolors.RED + '[+]' + globals.bcolors.WHITE + " A newer version is available: " + new_maldb_ver + "!" + print globals.bcolors.RED + '[+]' + globals.bcolors.WHITE + " Updating..." + + # Get the new DB and update it + + self.download_from_repo(globals.vars.db_path) + self.db.close_connection() + remove(globals.vars.db_path) + rename("maldb.db", globals.vars.db_path) + self.db.renew_connection() # Write the new DB version into the file + f = open(globals.vars.maldb_ver_file, 'w') f.write(new_maldb_ver) f.close() - - # Get the new CSV and update it - csvurl = globals.vars.giturl_dl_dl + globals.vars.main_csv_file - u = urllib2.urlopen(csvurl) - f = open(globals.vars.main_csv_file, 'wb') - meta = u.info() - file_size = int(meta.getheaders("Content-Length")[0]) - print "Downloading: %s Bytes: %s" % (globals.vars.main_csv_file, file_size) - file_size_dl = 0 - block_sz = 8192 - while True: - buffer = u.read(block_sz) - if not buffer: - break - file_size_dl += len(buffer) - f.write(buffer) - status = r"%10d [%3.2f%%]" % ( - file_size_dl, file_size_dl * 100. / file_size) - status = status + chr(8) * (len(status) + 1) - print status, - f.close() + return def get_malware(self, id): + # get mal location - db = db_handler.DBHandler() - loc = db.query("SELECT LOCATION FROM MALWARES WHERE ID=?", id)[0][0] - print loc + + loc = self.db.query("SELECT LOCATION FROM MALWARES WHERE ID=?", id)[0][0] + + # get from git + self.download_from_repo(loc, '.zip') self.download_from_repo(loc, '.pass') self.download_from_repo(loc, '.md5') self.download_from_repo(loc, '.sha256') - # get from git - def download_from_repo(self, mal_location, suffix): + def download_from_repo(self, filepath, suffix=''): if globals.vars.DEBUG_LEVEL is 1: print locals() - file_name = mal_location.rsplit('/')[-1] + suffix - url = globals.vars.giturl_dl + mal_location + '/' + file_name + file_name = filepath.rsplit('/')[-1] + suffix + + # Dirty way to check if we're downloading a malware + + if suffix is not '': + url = globals.vars.giturl_dl + filepath + '/' + file_name + else: + url = globals.vars.giturl_dl + filepath u = urllib2.urlopen(url) f = open(file_name, 'wb') meta = u.info() @@ -117,3 +114,4 @@ class Updater: status = status + chr(8) * (len(status) + 1) sys.stdout.write('\r' + status) f.close() + print "\n" diff --git a/theZoo.py b/theZoo.py index 69c71ec..29ed640 100644 --- a/theZoo.py +++ b/theZoo.py @@ -20,7 +20,7 @@ import sys import os from optparse import OptionParser -from imports.updatehandler import Updater +from imports.update_handler import Updater from imports import manysearches from imports import muchmuchstrings from imports.eula_handler import EULA