update.py

#!/usr/bin/env python
# This file is part of Viper - https://github.com/botherder/viper
# See the file 'LICENSE' for copying permission.

import os
import hashlib
from StringIO import StringIO
from zipfile import ZipFile, ZIP_STORED

from viper.common.out import *
from viper.common.network import download
from viper.common.objects import File
from viper.common.utils import path_split_all

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

# 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 main():
    print_warning("WARNING: If you proceed you will lose any changes you might have made to Viper.")
    choice = raw_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 StringIO, we will store the master.zip data in here.
    zip_data = StringIO()
    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()

    # 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}.
        local_file_path = os.path.join(*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.
        new_local = open(local_file_path, 'w')
        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()

if __name__ == '__main__':
    main()