viper-update

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This file is part of Viper - https://github.com/viper-framework/viper
# See the file 'LICENSE' for copying permission.

import os
import sys
import hashlib
from zipfile import ZipFile
from optparse import OptionParser
from sqlalchemy import create_engine
from datetime import datetime

try:
    from io import BytesIO
except ImportError:
    from BytesIO import BytesIO

from viper.common.out import print_info
from viper.common.out import print_warning
from viper.common.out import print_error
from viper.common.out import print_success
from viper.common.out import print_item
from viper.common.network import download
from viper.common.objects import File
from viper.core.config import __config__
from viper.core.project import __project__

# For python2 & 3 compatibility, a bit dirty,
# but it seems to be the least bad one.
try:
    input = raw_input
except NameError:
    pass

url = 'https://github.com/viper-framework/viper/archive/master.zip'


# Taken from the Python Cookbook.
def path_split_all(path):
    allparts = []
    while 1:
        parts = os.path.split(path)
        if parts[0] == path:
            allparts.insert(0, parts[0])
            break
        elif parts[1] == path:
            allparts.insert(0, parts[1])
            break
        else:
            path = parts[0]
            allparts.insert(0, parts[1])

    return allparts


# TODO: this is a first draft, needs more work.
# - Add a check for current working directory.
# - Add error handling.
# - Ignore all git related files and directories.
def update():
    print_warning("WARNING: If you proceed you will lose any changes you might have made to Viper.")
    choice = input("Are you sure you want to proceed? [y/N] ")

    if choice.lower() != 'y':
        return

    # Download the latest Zip archive from GitHub's master branch.
    master = download(url)
    # Instantiate a BytesIO, we will store the master.zip data in here.
    zip_data = BytesIO()
    zip_data.write(master)
    # Initialize the Zip archive.
    zip_file = ZipFile(zip_data, 'r')
    # Obtain a list of all the files contained in the master.zip archive.
    names = zip_file.namelist()

    # Get appropriate root folder.
    cfg = __config__
    if cfg.paths.storage_path:
        base_path = cfg.paths.storage_path
        resources_path = cfg.paths.storage_path
    else:
        base_path = os.path.join(os.getenv('HOME'), '.viper')
        resources_path = '/usr/local/share/viper'

    # Loop through all file and directories in master.zip.
    for name in names[1:]:
        # Split the path in parts.
        name_parts = path_split_all(name)
        # We strip the base directory, which is generated by GitHub in the
        # master.zip archive as {project}-{branch}.
        # Also, if a folder is stored in 'data', means that it is static
        # resources we embedded. In case of global install, we want to
        # install it in share folder.
        if name_parts[1] == 'data':
            local_file_path = os.path.join(resources_path, *name_parts[1:])
        else:
            local_file_path = os.path.join(base_path, *name_parts[1:])

        # Skip if the entry is a directory.
        if os.path.isdir(local_file_path):
            continue

        # Read the data of the current file.
        name_data = zip_file.read(name)
        # Calculate MD5 hash of the new file.
        name_data_md5 = hashlib.md5(name_data).hexdigest()

        # If the file already exists locally, we check if its MD5 hash
        # matches the one of the newly downloaded copy. If it does, we
        # obviously skip it.
        exists = False
        if os.path.exists(local_file_path):
            exists = True
            if File(local_file_path).md5 == name_data_md5:
                print_info("{0} up-to-date".format(local_file_path))
                continue

        # Open the local file, whether it exists or not, and either
        # rewrite or write the new content.
        try:
            new_local = open(local_file_path, 'wb')
        except IOError as e:
            if e.errno == 21:
                # It's a new directory.
                try:
                    os.mkdir(local_file_path)
                except Exception as e:
                    print_error("Uname to create new directory {0}: {1}".format(local_file_path, e))
                else:
                    print_success("New directory {0} has been created".format(local_file_path))

                continue

        new_local.write(name_data)
        new_local.close()

        if exists:
            print_success("File {0} has been updated".format(local_file_path))
        else:
            print_success("New file {0} has been created".format(local_file_path))

    zip_file.close()
    zip_data.close()


def update_db():
    print_item("Backing up Sqlite DB")

    # backup of database name with a timestamp to avoid to be overwritten
    db_backupname = "viper_{0}.db.bak".format(datetime.utcnow().strftime("%Y%m%d-%H%M%S"))

    try:
        os.rename('viper.db', db_backupname)
    except Exception as e:
        print_error("Failed to Backup. {0} Stopping".format(e))
        return

    print_item("Creating New DataBase File")
    from viper.core.database import Database
    Database()

    print_item("Connecting to Viper Databases")
    old_engine = create_engine('sqlite:///{0}'.format(db_backupname))
    db_path = os.path.join(__project__.get_path(), 'viper.db')
    new_engine = create_engine('sqlite:///{0}'.format(db_path))

    print_item("Reading data from Old Database")
    malware = old_engine.execute('SELECT * FROM malware').fetchall()
    association = old_engine.execute('SELECT * FROM association').fetchall()
    notes = old_engine.execute('SELECT * FROM note').fetchall()
    tags = old_engine.execute('SELECT * FROM tag').fetchall()

    print_item(" Adding rows to New Database")

    # Add all the rows back in
    for row in notes:
        new_engine.execute("INSERT INTO note VALUES ('{0}', '{1}', '{2}')".format(row[0], row[1], row[2]))

    for row in tags:
        new_engine.execute("INSERT INTO tag VALUES ('{0}', '{1}')".format(row[0], row[1]))

    for row in malware:
        new_engine.execute("INSERT INTO malware VALUES ('{0}', '{1}', '{2}', '{3}', '{4}', "
                           "'{5}', '{6}', '{7}', '{8}', '{9}', '{10}', '{11}', 'Null')".format(row[0], row[1], row[2],
                                                                                               row[3], row[4], row[5],
                                                                                               row[6], row[7], row[8],
                                                                                               row[9], row[10], row[11])
                           )

    # Rebuild association table with foreign keys
    for row in association:
        if row[0] is None:
            tag_id = "Null"
        else:
            tag_id = "(SELECT id from tag WHERE id='{0}')".format(row[0])
        if row[1] is None:
            note_id = "Null"
        else:
            note_id = "(SELECT id from note WHERE id='{0}')".format(row[1])
        if row[2] is None:
            malware_id = "Null"
        else:
            malware_id = "(SELECT id from malware WHERE id='{0}')".format(row[2])

        new_engine.execute("INSERT INTO association VALUES ({0}, {1}, {2}, 'Null')".format(tag_id, note_id, malware_id))

    print_info("Update Complete")


if __name__ == '__main__':
    parser = OptionParser(usage='usage: %prog -c|-d')
    parser.add_option("-d", "--db", action='store_true', default=False, help="Update DB Tables")
    parser.add_option("-c", "--core", action='store_true', default=False, help="Update Core Files")

    (options, args) = parser.parse_args()

    if not options.db and not options.core:
        print("")
        print("=========================================================================")
        print("| This script will Update your Viper to the latest Dev Version.         |")
        print("| Some Dev Updates have new or modified tables. This script will allow  |")
        print("| you to update your existing DB files to the required Schema.          |")
        print("=========================================================================\n")
        parser.print_help()
        sys.exit()

    if options.db:
        print_warning("To update Projects you will need to copy their viper.db file in to the main viper folder")
        print_warning("Run the DB update then move the new db file back to the project folder. ")
        print_info("Updating to New DB format")

        update_db()

    if options.core:
        update()