#! /usr/bin/env python

import os
import sys
import tuf
from tuf.client import updater
import tuf.conf
import optparse
import shutil
import hashlib

def listdir(dir):
    output = set([])
    dirlist = os.listdir(dir)
    for file in dirlist:
        if os.path.isfile(os.path.join(dir,file)):
            output.add(os.path.join(dir,file))
        if os.path.isdir(os.path.join(dir,file)):
            output.update(listdir(os.path.join(dir,file)))
    return output
    

def pull_directory(source):
    """copies source into .files"""

    shutil.copytree(source, os.path.join(os.getcwd(),'.files'))

def run_tuf():
    """Updates all of the files in ./.files/ dir."""

    print "Downloading Files"

    tuf.conf.settings.repo_meta_dir = "./"

    repo = updater.Repository(".nmap", {'repo': {'urlbase': 'http://updates-test.nmap.org/server_root', 'metapath': "meta", 'targetspath': "targets", 'metacontent': ['**'], 'targetscontent': ['**']}})

    #update
    repo.refresh()

    targets = repo.get_all_targets()
    files_to_update = set(repo.get_files_to_update(targets))
    
    downloaded = 0
    for target in targets:
        if target not in files_to_update:
            continue
        path = target.path
        dir_path = os.path.dirname(path)
        if dir_path != '':
            try:
                os.makedirs(dir_path)
            except OSError, e:
                pass
        print "Downloading " + target.path
        target.download(target.path)
        downloaded +=1

    print "Downloaded: " + str(downloaded) + " Files"

    initial_files = listdir('./.files')
    repo.remove_missing_files()
    new_files = listdir('./.files')

    deleted_files = initial_files-new_files

    return deleted_files

def copy(source, destination):
    dir_path = os.path.dirname(destination)
    if dir_path != '':
        try:
            os.makedirs(dir_path)
        except OSError, e:
            pass
    shutil.copy2(source, destination)

def push_directory(destination, options, deleted_files):
    "Moves .files into destination"

    if destination != None:
        destination = [ os.path.join(destination, "bin"),
                    os.path.join(destination, "lib"),
                    os.path.join(destination, "share")]
    else:
        destination = [None,None,None]

    if options.bindir != None:
        destination[0] = options.bindir
    if options.libdir != None:
        destination[1] = options.libdir
    if options.sharedir != None:
        destination[2] = options.sharedir    

    source = [  os.path.join(os.getcwd(),'.files/bin/'),
                os.path.join(os.getcwd(),'.files/lib/'),
                os.path.join(os.getcwd(),'.files/share/')]

    for i in range(0,3):
        print 'Installing: ' + destination[i]
        dirs = os.walk(source[i])
        for info in dirs:
            for file_name in info[2]:
                source_path = os.path.join(info[0],file_name)
                destination_path = source_path[len(source[i]):]
                destination_path = os.path.join(destination[i] , destination_path)
                copy(source_path,destination_path)
    
    for path in deleted_files:
        for i in range(0,3):
            if path.find(source[i]) >= 0:
                path.replace(source[i],destination[i])
                if os.path.exists(path):
                    #os.remove(pathname)
                    print 'Removed ' + str(path)
                continue


def run(options):
    all_specified = options.bindir != None and options.libdir != None \
                                         and options.sharedir != None

    if not (options.source or options.prefix or all_specified):
        print 'Please Specify a prefix with --prefix'

    if options.source:
        pull_directory(options.source)

    if options.prefix or all_specified:
        deleted_files = run_tuf()
        push_directory(options.prefix, options, deleted_files)

if __name__ == "__main__":
    tuf.log.set_log_level("ERROR")

    #Get User Arguments
    parser = optparse.OptionParser(description='Update Nmap')
    parser.add_option('--prefix', type=str,default=None, help='Destination for files')
    parser.add_option('--source', type=str,default = None, help='Source for files. Should only be set on the server')
    parser.add_option('--debug', action= 'store_true', default = False, help='Do not catch any exceptions')
    parser.add_option('--bindir', type=str, default=None, help='custom bin dir')
    parser.add_option('--libdir', type=str, default = None, help='custom lib dir')
    parser.add_option('--sharedir', type=str, default=None, help='custom share dir')
    args, rest = parser.parse_args()

    if args.debug == True:
        run(args)
        sys.exit(1)

    try:
        run(args)
    except os.error, e:
        print "Error:" , e
    except KeyboardInterrupt, e:
        print "Aborting"
    except IOError, e:
        print "Could not access files, did you try to install without root permissions?"
    except (tuf.DownloadError, tuf.RepoError), e:
        print "DOWNLOAD ERROR. The server may be compromised or set up incorrectly"
    except (tuf.BadHash, tuf.BadSignature, tuf.BadPassword, tuf.UnknownKeyError), e:
        print "AUTHENTICATION ERROR. The client or server may be compromised."
    except (tuf.FormatException), e:
        print "FORMAT EXCEPTION. The server may be compromised."
    except (tuf.UnsupportedAlgorithmError, tuf.UnknownFormat), e:
        print "UNSUPPORTED. Client may be old or server may be compromised."
    except (tuf.InternalError, tuf.MetadataNotAvailableError, tuf.UnknownMethod), e:
        print "TUF ERROR. Server may be compromised."
    except (tuf.CryptoError, tuf.PubkeyFormatException, tuf.ExpiredMetadataError), e:
        print "CRYPTO ERROR. Server may be compromised"
    except (tuf.CheckNotSupported, tuf.KeyAlreadyExistsError, tuf.RoleAlreadyExistsError, tuf.UnknownRoleError, tuf.InvalidNameError), e:
        print "SERVERSIDE ERROR. Server may be compromised"