# Hello, this script is written in Python - see http://python.org
"""d_avpup.py 1.0.1 - AVP signatures files auto-download.

This Python script is a handy replacement of the AVP auto-update feature.

This script:
        - will query Kaspersky site for updates
        - will download necessary ZIP files
        - will read AVP/KAV signatures path from the registry
        - will unzip the files to this path.

ZIP files will be downloaded in the current directory (the script directory).
ZIP files already downloaded will not be downloaded again (except daily.zip).

To force this script to download everything again, delete all ZIP files from
current directory.

Tested with AVP 3.5.  Should work with other versions.

License :
    This script is public domain.
    
Author :
    Sebastien SAUVAGE <sebsauvage at sebsauvage dot net>
    http://sebsauvage.net
"""
import sys, urllib, re, zipfile, os, os.path

avp_update_url     = 'http://www.kaspersky.com/bases/' 
#os.environ['http_proxy'] = ''   # Enter address of your proxy if any (eg. 'http://proxy.free.fr')

try:
    import win32api, win32con
except:
    print "This script requires the win32all extensions for Python."
    print "See http://starship.python.net/crew/mhammond/win32/"
    print " or http://www.python.org/windows/win32/"
    sys.exit(1)

def getbasepath():
    """Locate AVP shared files directory by reading regitry key:
       HKEY_LOCAL_MACHINE\SOFTWARE\KasperskyLab\SharedFiles\ VEDataFilePath
       (Signatures are in the ..\bases subdirectory)"""
    try: # Open registry key
        regkey = win32api.RegOpenKeyEx( win32con.HKEY_LOCAL_MACHINE, 'SOFTWARE\\KasperskyLab\\SharedFiles\\', 0, win32con.KEY_READ )
    except:
        raise RuntimeError, 'Could not open regitry key HKEY_LOCAL_MACHINE\\SOFTWARE\\KasperskyLab\\SharedFiles\\'

    try: # Read regitry key
        (vepath,keytype) = win32api.RegQueryValueEx(regkey, 'VEDataFilePath' )
    except:
        raise RuntimeError, 'Could not read key VEDataFilePath from registry'

    if vepath[-1] != '\\': vepath += '\\' 
    return vepath + 'bases'   # the path to the signature files (*.avc files)

def download_bases( avpurl ):
    """Download AVP signature files in the current directory.
       Returns a list of downloaded files."""
    if avpurl[-1] != '/' : avpurl += '/'     
    print 'Fecthing update page from %s' % avpurl
    page = urllib.urlopen(avpurl).read()
    # Now extract all <A HREF> tags pointing to ZIP files (eg. '<A HREF="daily.zip">')
    matches = re.compile(r'<a href="(.*?\.zip)">', re.IGNORECASE | re.MULTILINE ).findall(page)
    if not matches:
        print 'Could not understand page %s.' % avpurl
        sys.exit(1)
    print 'Files to download :\n%s\n' % ','.join(matches)
    # Download all files, except the one we already have
    for filename in matches:
        if os.path.isfile(filename) and filename.lower() != 'daily.zip':  # (Always download daily.zip)
            print '%s already downloaded. Skipping.' % filename
        else:
            sys.stdout.write( 'Downloading %s ...' % filename )
            file = open(filename,'w+b')
            file.write(  urllib.urlopen(avpurl + filename).read() )
            file.close()
            sys.stdout.write(' done.\n')
    return matches   

def unziptobasepath( zipfilename, path ):
    """Unzips filename to path (overwriting files if present)."""
    sys.stdout.write( "Unzipping %s : " % zipfilename )
    zfile = zipfile.ZipFile(zipfilename,'r')
    for filename in zfile.namelist():
        sys.stdout.write('#')
        file = open( os.path.join(path,filename), 'w+b')
        file.write( zfile.read(filename) )
        file.close()
    sys.stdout.write('\n')
    
print 'AVP auto-update 1.0.1 - by Sebastien SAUVAGE <sebsauvage at sebsauvage dot net>'
basepath = getbasepath()
files = download_bases(avp_update_url)
print ''

files.sort()

print 'Unzipping files to %s' % basepath
# First unzip cumulative update avp*.zip
for filename in files:
    if filename[:3].lower() == 'avp':
        unziptobasepath( filename, basepath )

# then unzip weekly updates up*.zip
for filename in files:
    if filename[:2].lower() == 'up':
        unziptobasepath( filename, basepath )

# then unzip daily update daily.zip
for filename in files:
    if filename.lower() == 'daily.zip':
        unziptobasepath( filename, basepath )

print 'All done.'