'''Various utility functions for handling Debian-like archive structures.

Copyright (C) 2006 Canonical Ltd.
Author: Martin Pitt <martin.pitt@ubuntu.com>

This program is free software; you can redistribute it and/or modify it
under the terms of the GNU General Public License as published by the
Free Software Foundation; either version 2 of the License, or (at your
option) any later version.  See http://www.gnu.org/copyleft/gpl.html for
the full text of the license.
'''

import apt_pkg, urllib, os, gzip, tempfile, os.path, shutil, tarfile, time
import subprocess, glob
from cStringIO import StringIO

apt_pkg.InitSystem()

def debcontrol_map(url, map_key='Package'):
    '''Read a debcontrol format file from the given URL and parse it into a map 
    map_key->key->value.

    map_key defines a 'primary key' of the possible keys in the file.
    E. g. debcontrol_map('Sources.gz')['dpkg']['Version'] would deliver the
    Version of the dpkg source package.'''

    # write the uncompressed downloaded list into a temporary file
    downloadfile = urllib.urlretrieve(url)[0]
    list = tempfile.TemporaryFile()
    try:
        print >> list, gzip.open(downloadfile).read()
    except IOError:
        print >> list, open(downloadfile).read()
    list.seek(0)

    map = {}
    parser = apt_pkg.ParseTagFile(list)
    while parser.Step() == True:
        pkg = {}
        for key in parser.Section.keys():
            if key != map_key:
                pkg[key] = parser.Section.get(key)
        map[parser.Section.get(map_key)] = pkg
    return map

def debcontrol_dominate(infile, outfile, group_by='Package', 
    sort_key='Version', cmp=apt_pkg.VersionCompare):
    '''Read a file-like object in debcontrol syntax, group the records by the
    key group_by and only retain the member with the largest sort_key value
    per group.
    
    The result is written to 'outfile' (which must be a file-like object). The
    comparison defaults to apt_pkg.VersionCompare (since this should be the
    main use case, but another signature-compatible comparison function may be
    specified.'''

    max = {} # group_by -> (sort_key_value, text)
    parser = apt_pkg.ParseTagFile(infile)
    while parser.Step() == True:
	if cmp(max.get(parser.Section[group_by], ('', ''))[0], parser.Section[sort_key]) < 0:
	    max[parser.Section[group_by]] = (parser.Section[sort_key], str(parser.Section))

    ks = max.keys()
    ks.sort()
    first = True
    for k in ks:
	if first:
	    first = False
	else:
	    outfile.write('\n')
	outfile.write(max[k][1])

def install_pool(archive_root, file_url, component, source):
    '''Install given file into the pool, with the given component and source
    package name.
    
    Return True on success, or False if the target file already exists.'''

    if source.startswith('lib'):
	hash = source[:4]
    else:
	hash = source[0]

    destdir = os.path.join(archive_root, 'pool', component, hash, source)
    destfile = os.path.join(destdir, os.path.basename(file_url))
    if os.path.exists(destfile):
	return False
    if not os.path.isdir(destdir):
	os.makedirs(destdir)

    urllib.urlretrieve(file_url, destfile)

    return True

def tar_install_pool(archive_root, tar_url, component, source):
    '''Unpack the given tarball URL and install its files into the pool,
    belonging to the given component and source package name.

    Return True on success, or False if the target file already exists.'''

    tar = urllib.urlretrieve(tar_url)[0]
    assert tarfile.is_tarfile(tar)

    if source.startswith('lib'):
	hash = source[:4]
    else:
	hash = source[0]

    destdir = os.path.join(archive_root, 'pool', component, hash, source)
    if not os.path.isdir(destdir):
	os.makedirs(destdir)

    t = tarfile.open(tar)
    for m in t.getmembers():
	if os.path.exists(os.path.join(destdir, m.name)):
	    return False
	t.extract(m, destdir)

    if tar != tar_url:
        # clean up temporary file created by urlretrieve()
        os.unlink(tar)

    return True

def create_indexes(archive_root, dist, components, archs, ext='.deb',
    stderr=subprocess.PIPE, ref_archive=None, gpg_sign=False, gpg_keyid=None):
    '''Create an apt-ftparchive configuration file, create the necessary
    directory structures and all apt-ftparchive to create Packages.gz files.
    
    ext specifies the file name extension (defaulting to '.deb').
    stderr allows you to redirect the apt-ftparchive output (e. g.
    None to use the caller's stderr instead of ignoring it).
    
    If ref_archive is given, it needs to point to an existing archive URL.
    create_indexes() will then download the Packages/Sources indexes from that
    one and only consider the package versions mentioned there.

    If gpg_sign is True, "gpg -abs --batch -o Release.gpg Release" will be
    called on Release files. If gpg_keyid is specified, it is passed as
    --default-key.
    '''
    # create configuration file
    conf = tempfile.NamedTemporaryFile()
    print >> conf, '''Dir {
	ArchiveDir ".";
	FileListDir "./lists";
}

Default {
	Packages::Extensions "%s";
}

TreeDefault {
	Directory "./pool";
	SrcDirectory "./pool";
	BinCacheDB "./db";
	Contents " ";
	FileList "$(SECTION)-$(ARCH)";
}

Tree "dists/%s" {
	Sections "%s";
	Architectures "%s";
}''' % (ext, dist, components, archs)
    conf.flush()

    old_cwd = os.getcwd()
    os.chdir(archive_root)
    try:
	# create directories and file lists
	if not os.path.isdir('lists'):
	    os.mkdir('lists')
	for component in components.split():
	    pooldir = os.path.join('pool', component)
	    if not os.path.isdir(pooldir):
		os.makedirs(pooldir)
	    for arch in archs.split():
		if arch == 'source':
		    d = os.path.join('dists', dist, component, 'source')
		else:
		    d = os.path.join('dists', dist, component, 'binary-' + arch)
		if not os.path.isdir(d):
		    os.makedirs(d)
                if ref_archive:
                    # create a package -> version map for this
                    # component/architecture
                    if arch == 'source':
                        url = os.path.join(ref_archive, 'dists', dist, component, 'source', 'Sources.gz')
                    else:
                        url = os.path.join(ref_archive, 'dists', dist, component, 'binary-' + arch, 'Packages.gz')
                    ref_map = debcontrol_map(url)
                    listfile = open(os.path.join('lists', '%s-%s' %
                        (component, arch)), 'w')
                    for pkg, info in ref_map.iteritems():
                        if not pkg:
                            continue
                        ver = info['Version']
                        epoch_offset = ver.find(':')
                        if epoch_offset > 0:
                            ver = ver[epoch_offset+1:]
                        source = info.get('Source', pkg)
                        if source.startswith('lib'):
                            srcprefix = source[:4]
                        else:
                            srcprefix = source[0]
                        if arch == 'source':
                            raise SystemError, 'create_index() does not support source packages with a reference_archive yet'
                        else:
                            listed = set()
                            for f in glob.glob(os.path.join(pooldir, srcprefix,
                                source, '%s*_%s_%s%s' % (pkg, ver, arch, ext))):
                                if not f in listed:
                                    print >> listfile, f
                                    listed.add(f)
                            listed = None
                else:
                    # just use everything we find
                    subprocess.call("find %s -name '*_%s%s' > lists/%s-%s" %
                        (pooldir, arch, ext, component, arch), shell=True)

	# call apt-ftparchive
	assert subprocess.call(['apt-ftparchive', 'generate', conf.name],
	    stderr=stderr) == 0
        os.chdir(os.path.join('dists', dist))
        apt_ftparchive = subprocess.Popen(['apt-ftparchive', '-o',
            'APT::FTPArchive::Release::Suite='+dist, 'release', '.'],
	    stdout=subprocess.PIPE, stderr=stderr)
        out = apt_ftparchive.communicate()[0]
        assert apt_ftparchive.returncode == 0
	r = open('Release', 'w')
        r.write(out)
        r.close()
        if gpg_sign:
            try:
                os.unlink('Release.gpg')
            except OSError:
                pass
            gpg = ['gpg', '-abs', '--batch']
            if gpg_keyid:
                gpg += ['--default-key', gpg_keyid]
            gpg += ['-o', 'Release.gpg', 'Release']
            assert subprocess.call(gpg, stdout=subprocess.PIPE, stderr=stderr) == 0
    finally:
	os.chdir(old_cwd)
	if os.path.isdir(os.path.join(archive_root, 'lists')):
	    shutil.rmtree(os.path.join(archive_root, 'lists'))

def _packages_files(path, files):
    '''Add all files from given Packages path to the 'files' set.'''

    parser = apt_pkg.ParseTagFile(open(path))
    while parser.Step() == True:
	f = parser.Section['Filename']
	if f.startswith('./'):
	    f = f[2:]
	files.add(f)

def _sources_files(path, files):
    '''Add all files from given Sources path to the 'files' set.'''

    parser = apt_pkg.ParseTagFile(open(path))
    while parser.Step() == True:
	d = parser.Section.get('Directory')
	for l in parser.Section['Files'].splitlines():
	    fields = l.split()
	    assert len(fields) == 3
	    if d:
		files.add(os.path.join(d, fields[2]))
	    else:
		files.add(fields[2])

def archive_cleanup(archive_root, keep_days=0):
    '''Remove all files from an archive pool that are not mentioned in any
    Packages or Sources file and which are older than keep_days.'''

    old_cwd = os.getcwd()
    os.chdir(archive_root)

    try:
	# determine set of all indexed files (which we want to keep)
	distroot = 'dists'
	if not os.path.isdir(distroot):
	    distroot = '.'

	indexed_files = set([])
	for root, dirs, files in os.walk(distroot):
	    if root == '.':
		root = ''
	    elif root.startswith('./'):
		root = root[2:]
	    if 'Packages' in files:
		_packages_files(os.path.join(root, 'Packages'), indexed_files)
		indexed_files.add(os.path.join(root, 'Packages'))
		indexed_files.add(os.path.join(root, 'Packages.gz'))
		indexed_files.add(os.path.join(root, 'Packages.bz2'))
	    else:
		assert 'Packages.gz' not in files
		assert 'Packages.bz2' not in files

	    if 'Sources' in files:
		_sources_files(os.path.join(root, 'Sources'), indexed_files)
		indexed_files.add(os.path.join(root, 'Sources'))
		indexed_files.add(os.path.join(root, 'Sources.gz'))
		indexed_files.add(os.path.join(root, 'Sources.bz2'))
	    else:
		assert 'Sources.gz' not in files
		assert 'Sources.bz2' not in files

        now = time.time()
        keep_seconds = keep_days * 24 *3600

	# now walk the pool and clean up
	poolroot = 'pool'
	if not os.path.isdir(poolroot):
	    poolroot = '.'
	for root, dirs, files in os.walk(poolroot):
	    if root == '.':
		root = ''
	    if root.startswith('./'):
		root = root[2:]
	    for f in files:
		n = os.path.join(root, f) 
		if n not in indexed_files:
                    if keep_days:
                        if (now - os.stat(n).st_mtime) <= keep_seconds:
                            continue
		    os.unlink(n)
    finally:
	os.chdir(old_cwd)

def _pack(src, dest):
    '''Pack uncompressed file src to compressed file dest.

    dest must end with .gz for gzip compression or .bz2 for bzip2
    compression. This uses a temporary file so that the resulting compressed
    file is not clobbered with a half-written new version.'''

    if dest.endswith('.gz'):
	z = gzip.open(dest + '.new', 'w', 9)
    elif dest.endswith('.bz2'):
	z = bz2.BZ2File(dest + '.new', 'w')
    f = open(src)
    while True:
	block = f.read(10240)
	if not block:
	    break
	z.write(block)
    f.close()
    z.close()
    os.rename(dest + '.new', dest)

def archive_dominate(archive_root, keep_days=0):
    '''Dominate all Sources and Packages files, refresh the corresponding *.gz
    files, and call archive_cleanup().
    
    keep_days is passed to archive_cleanup(), i. e. all files which are younger
    than keep_days will be retained.'''

    # determine set of all indexed files (which we want to keep)
    distroot = os.path.join(archive_root, 'dists')
    if not os.path.isdir(distroot):
	distroot = archive_root

    for root, dirs, files in os.walk(distroot):
	if 'Packages' in files:
	    path = os.path.join(root, 'Packages')
	    debcontrol_dominate(open(path), open(path + '.new', 'w'))

	    # regenerate compressed files
	    if 'Packages.gz' in files: 
		_pack(path + '.new', path + '.gz')
	    if 'Packages.bz2' in files: 
		_pack(path + '.new', path + '.bz2')
	    os.rename(path + '.new', path)
	else:
	    assert 'Packages.gz' not in files
	    assert 'Packages.bz2' not in files

	if 'Sources' in files:
	    path = os.path.join(root, 'Sources')
	    debcontrol_dominate(open(path), open(path + '.new', 'w'))

	    # regenerate compressed files
	    if 'Sources.gz' in files: 
		_pack(path + '.new', path + '.gz')
	    if 'Sources.bz2' in files: 
		_pack(path + '.new', path + '.bz2')
	    os.rename(path + '.new', path)
	else:
	    assert 'Sources.gz' not in files
	    assert 'Sources.bz2' not in files

    archive_cleanup(archive_root, keep_days)

#
# Test suite
#

import unittest, sys, signal
import SimpleHTTPServer, BaseHTTPServer

class _DebcontrolMapTest(unittest.TestCase):
    '''debcontrol_map()'''

    def setUp(self):

        self.tmpfiles = set(os.listdir('/tmp'))

	self.sources = tempfile.NamedTemporaryFile()
	print >> self.sources, '''Package: aalib
Binary: libaa-bin, libaa1-dev, libaa1
Version: 1.4p5-30
Directory: pool/main/a/aalib
Files:
 ad2e6b4821fb03fa6de7371a1b292f7b 733 aalib_1.4p5-30.dsc
 9801095c42bba12edebd1902bcf0a990 391028 aalib_1.4p5.orig.tar.gz

Package: abiword
Binary: abiword-plugins, abiword-plugins-gnome, abiword-help, abiword, abiword-common, abiword-gnome
Version: 2.4.4-0ubuntu5'''
	self.sources.flush()

	self.packages = tempfile.NamedTemporaryFile()
	print >> self.packages, '''Package: abiword
Architecture: amd64
Version: 2.4.4-0ubuntu5
Description: Short
 Long
 Long

Package: abiword-common
Source: abiword
Version: 2.4.4-0ubuntu5'''
	self.packages.flush()

    def tearDown(self):
	self.sources.close()
	self.packages.close()

        self.assertEqual(set(os.listdir('/tmp')), self.tmpfiles, 
                'does not leave temporary files around')

    def test_uncompressed(self):
	'''debcontrol_map() for uncompressed files'''

	i = debcontrol_map('file://' + self.sources.name)
	self.assertEqual(_mapkeystr(i), 'aalib abiword')
	self.assertEqual(_mapkeystr(i['aalib']), 'Binary Directory Files Version')
	self.assertEqual(_mapkeystr(i['abiword']), 'Binary Version')
	self.assertEqual(i['aalib']['Version'], '1.4p5-30')
	self.assertEqual(i['abiword']['Version'], '2.4.4-0ubuntu5')

	i = debcontrol_map('file://' + self.packages.name)
	self.assertEqual(_mapkeystr(i), 'abiword abiword-common')
	self.assertEqual(_mapkeystr(i['abiword']), 'Architecture Description Version')
	self.assertEqual(_mapkeystr(i['abiword-common']), 'Source Version')
	self.assertEqual(i['abiword']['Version'], '2.4.4-0ubuntu5')
	self.assertEqual(i['abiword']['Description'], 'Short\n Long\n Long')

	self.assertRaises(KeyError, i.__getitem__, 'nonexistant')
	self.assertRaises(KeyError, i['abiword'].__getitem__, 'nonexistant')

    def test_map_key(self):
	'''debcontrol_map for an alternate mapping key'''

	i = debcontrol_map('file://' + self.sources.name, 'Version')
	self.assertEqual(_mapkeystr(i), '1.4p5-30 2.4.4-0ubuntu5')
	self.assertEqual(_mapkeystr(i['1.4p5-30']), 'Binary Directory Files Package')

    def test_nonexisting(self):
	'''debcontrol_map() with a nonexisting URL'''

	self.assertRaises(IOError, debcontrol_map, 'file:///notexisting')

def _ls(dir):
    '''Return an 'ls'-like string with a sorted space-separated list of a
    directory contents.'''

    l = os.listdir(dir)
    l.sort()
    return ' '.join(l)

def _mapkeystr(map):
    '''Return string with space-separated sorted list of keys of a
    dictionary.'''

    k = map.keys()
    k.sort()
    return ' '.join(k)

def _find_files(dir, strip_prefix=None):
    '''Return set of all files contained in given dir, and recursively, its
    subdirectories.
    
    If strip_prefix is given, it is removed from all file names.'''

    if strip_prefix:
	plen = len(strip_prefix)
    else:
	plen = 0
    result = set([])
    for root, dirs, files in os.walk(dir):
	for f in files:
	    result.add(os.path.join(root, f)[plen:])
    return result

class _ArchiveTests(unittest.TestCase):
    '''Test functions that deal with an archive.'''

    def create_deb(self, package, version, arch, dir, ext='.deb'):
	'''Create a dummy deb with given name, version, and architecture, and
	write it into the given directory.
	
	ext can specify a nonstandard file name suffix. Return the full path of
	the deb.'''

	d = tempfile.mkdtemp()
	os.mkdir(os.path.join(d, 'DEBIAN'))
	open(os.path.join(d, 'DEBIAN', 'control'), 'w').write('''Package: %s
Version: %s
Priority: optional
Section: devel
Architecture: %s
Maintainer: Joe <joe@joe.com>
Description: dummy package
 just a test dummy package
''' % (package, version, arch))

	debpath = os.path.join(dir, '%s_%s_%s%s' % (package, version, arch, ext))
	assert subprocess.call(['dpkg', '-b', d, debpath],
	    stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0

	shutil.rmtree(d)
	assert os.path.exists(debpath)
	return debpath

    def create_source(self, package, version, dir, native=False):
	'''Create a dummy source package with given name and version, and write
	it into the given directory.

	If 'native' is True, a native source package is created (dsc+tar.gz),
	otherwise a normal dsc/diff.gz/orig.tar.gz. Return a list of paths of
	all created files.'''

	dir = os.path.abspath(dir)
	d = tempfile.mkdtemp()
	old_cwd = os.getcwd()
	os.chdir(d)

	os.makedirs(os.path.join(package, 'debian'))
	open(os.path.join(package, 'README'), 'w').write('hello\n')
	open(os.path.join(package, 'debian', 'control'), 'w').write(
'''Source: %s
Maintainer: Joe <joe@ubuntu.com>

Package: %s
''' % (package, package))
	open(os.path.join(package, 'debian', 'changelog'), 'w').write(
'''%s (%s) dummy; urgency=low

  * Dummy package.

 -- Joe <joe@joe.com>  Fri, 22 Sep 2006 13:07:40 +0200
''' % (package, version))
	open(os.path.join(package, 'debian', 'rules'), 'w')

	if not native:
	    upstreamver = '-'.join(version.split('-')[:-1])
	    t = tarfile.open('%s_%s.orig.tar.gz' % (package, upstreamver), 'w:gz')
	    t.add(os.path.join(package, 'README'))
	    t.close()

	assert subprocess.call(['dpkg-source', '-b', package],
	    stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0

	shutil.rmtree(package)
	created_files = []
	for f in os.listdir('.'):
	    shutil.move(f, dir)
	    created_files.append(os.path.realpath(f))

	os.chdir(old_cwd)
	shutil.rmtree(d)
	return created_files

    def setUp(self):
        self.tmpfiles = set(os.listdir('/tmp'))

	self.archivedir = tempfile.mkdtemp('archivedir', 'archive_tools.py')
	self.testfiles_dir = tempfile.mkdtemp('testfiles', 'archive_tools.py')

	# create test tarballs
	d = tempfile.mkdtemp()
	old_cwd = os.getcwd()
	os.chdir(d)

	self.foo_tar = os.path.join(self.testfiles_dir, 'foo_1-1_i386_debs.tar')
	t = tarfile.open(self.foo_tar, 'w')
	t.add(self.create_deb('foo1', '1-1', 'i386', '.'))
	t.add(self.create_deb('foo2', '1-1', 'i386', '.'))
	t.close()

	self.bar_tar = os.path.join(self.testfiles_dir, 'bar_2-2ubuntu1_powerpc_debs.tar')
	t = tarfile.open(self.bar_tar, 'w')
	t.add(self.create_deb('bar1', '2-2ubuntu1', 'powerpc', '.'))
	t.add(self.create_deb('bar2', '2-2ubuntu1', 'powerpc', '.'))
	t.close()

	os.chdir(old_cwd)
	shutil.rmtree(d)

    def tearDown(self):
	shutil.rmtree(self.archivedir)
	shutil.rmtree(self.testfiles_dir)

        self.assertEqual(set(os.listdir('/tmp')), self.tmpfiles, 
                'does not leave temporary files around')

    def test_install_pool(self):
	'''install_pool()'''

	f = tempfile.NamedTemporaryFile()
	print >> f, 'hello'
	f.flush()

	self.assert_(install_pool(self.archivedir, f.name, 'main', 'src1'))
	self.assert_(install_pool(self.archivedir, f.name, 'main', 'src2'))
	self.assert_(install_pool(self.archivedir, f.name, 'main', 'libsrc'))
	self.assert_(install_pool(self.archivedir, 'file://' + f.name, 'universe', 'src'))
	self.assert_(not install_pool(self.archivedir, f.name, 'universe', 'src'))

	self.assertEqual(_ls(os.path.join(self.archivedir, 'pool')), 'main universe')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'pool', 'main')), 'libs s')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'pool', 'main',
	    's')), 'src1 src2')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'pool', 'main',
	    'libs')), 'libsrc')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'pool', 'main',
	    'libs', 'libsrc')), os.path.basename(f.name))
	self.assertEqual(open(os.path.join(self.archivedir, 'pool', 'main',
	    'libs', 'libsrc', os.path.basename(f.name))).read(), 'hello\n')

	f.close()

    def test_tar_install_pool(self):
	'''tar_install_pool()'''

	# call tar_install_pool()
	self.assert_(tar_install_pool(self.archivedir, self.foo_tar, 'main', 'foo'))
	self.assertEqual(_ls(os.path.join(self.archivedir, 'pool', 'main', 'f',
	    'foo')), 'foo1_1-1_i386.deb foo2_1-1_i386.deb')

	p  = subprocess.Popen(['dpkg', '-I', os.path.join(self.archivedir, 'pool',
	    'main', 'f', 'foo', 'foo1_1-1_i386.deb')], stdout=subprocess.PIPE)
	self.assert_(p.communicate()[0].find('Package: foo1\n') >= 0)
	self.assertEqual(p.returncode, 0)

	p  = subprocess.Popen(['dpkg', '-I', os.path.join(self.archivedir, 'pool',
	    'main', 'f', 'foo', 'foo2_1-1_i386.deb')], stdout=subprocess.PIPE)
	self.assert_(p.communicate()[0].find('Package: foo2\n') >= 0)
	self.assertEqual(p.returncode, 0)

	# files exist now, should return False
	self.assert_(not tar_install_pool(self.archivedir, self.foo_tar, 'main', 'foo'))

    def _fork_http_server(self, dir):
        '''fork a HTTP server.'''

        # fork test HTTP server
        pid = os.fork()
        if pid == 0:
            try:
                os.chdir(dir)
                httpd = BaseHTTPServer.HTTPServer(('', 8427),
                    SimpleHTTPServer.SimpleHTTPRequestHandler)
                httpd.serve_forever()
                os._exit(0)
            except:
                print '********** HTTP server failed: ***********'
                traceback.print_exc()
                os._exit(1)
        else:
            self.http_pid = pid
            # wait until server is ready
            while True:
                time.sleep(0.1)
                try:
                    f = urllib.urlopen('http://localhost:8427')
                except IOError:
                    continue
                f.read()
                f.close()
                break

    def _join_http_server(self):
        '''Wait for the forked HTTP server to finish and assert that it exits with zero.'''

        os.kill(self.http_pid, signal.SIGTERM)
        os.wait()

    def test_tar_install_pool_http(self):
        '''tar_install_pool() with a http:// URL'''

        try:
            self._fork_http_server(os.path.dirname(self.foo_tar))

            self.assert_(tar_install_pool(self.archivedir, 'http://localhost:8427/' +
                os.path.basename(self.foo_tar), 'main', 'foo'))
            self.assertEqual(_ls(os.path.join(self.archivedir, 'pool', 'main', 'f',
                'foo')), 'foo1_1-1_i386.deb foo2_1-1_i386.deb')
        finally:
            self._join_http_server()

    def test_create_indexes(self):
	'''create_indexes()'''

	self.assert_(tar_install_pool(self.archivedir, self.foo_tar, 'universe', 'foo'))
	self.assert_(tar_install_pool(self.archivedir, self.bar_tar, 'main', 'bar'))
	create_indexes(self.archivedir, 'edgy', 'main restricted universe', 
	    'i386 powerpc amd64', stderr=subprocess.PIPE)

	# check directory structure
	self.assertEqual(_ls(self.archivedir), 'db dists pool')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy')),
            'Release main restricted universe')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy',
	    'main', 'binary-i386')), 'Packages Packages.gz')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy',
	    'universe', 'binary-powerpc')), 'Packages Packages.gz')

	# check Packages contents
	map = debcontrol_map(os.path.join(self.archivedir,
	    'dists', 'edgy', 'universe', 'binary-i386', 'Packages.gz'))
	self.assertEqual(map['foo1']['Architecture'], 'i386')
	self.assertEqual(map['foo1']['Version'], '1-1')
	self.assertEqual(map['foo1']['Filename'], 'pool/universe/f/foo/foo1_1-1_i386.deb')
	self.assert_(not map.has_key('bar1')) # since that's powerpc

	map = debcontrol_map(os.path.join(self.archivedir,
	    'dists', 'edgy', 'main', 'binary-powerpc', 'Packages.gz'))
	self.assertEqual(map['bar2']['Architecture'], 'powerpc')
	self.assertEqual(map['bar2']['Version'], '2-2ubuntu1')
	self.assertEqual(map['bar2']['Filename'], 'pool/main/b/bar/bar2_2-2ubuntu1_powerpc.deb')
	self.assert_(not map.has_key('foo1')) # since that's i386

	self.assertEqual(open(os.path.join(self.archivedir, 'dists', 'edgy', 'main',
	    'binary-i386', 'Packages')).read(), '')
	self.assertEqual(open(os.path.join(self.archivedir, 'dists', 'edgy', 'universe',
	    'binary-powerpc', 'Packages')).read(), '')
	self.assertEqual(open(os.path.join(self.archivedir, 'dists', 'edgy', 'main',
	    'binary-amd64', 'Packages')).read(), '')
	self.assertEqual(open(os.path.join(self.archivedir, 'dists', 'edgy', 'universe',
	    'binary-amd64', 'Packages')).read(), '')

        rel = open(os.path.join(self.archivedir, 'dists', 'edgy', 'Release')).read()
        self.assert_('\nSuite: edgy\n' in rel)
        self.failIf('dists/edgy' in rel)
        self.failIf('Release' in rel)
        self.assert_(' main/binary-amd64/Packages\n' in rel)

    def test_create_indexes_source(self):
	'''create_indexes() for source packages'''

	maindir = os.path.join(self.archivedir, 'pool', 'main')
	os.makedirs(os.path.join(maindir, 'f', 'foo'))
	self.create_source('foo', '1-1', os.path.join(maindir, 'f', 'foo'))
	os.makedirs(os.path.join(maindir, 'b', 'bar'))
	self.create_source('bar', '2-2', os.path.join(maindir, 'b', 'bar'), native=True)

	create_indexes(self.archivedir, 'edgy', 'main', 'source i386',
	    stderr=subprocess.PIPE)

	# check directory structure
	self.assertEqual(_ls(self.archivedir), 'db dists pool')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy')),
            'Release main')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy',
	    'main')), 'binary-i386 source')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy',
	    'main', 'binary-i386')), 'Packages Packages.gz')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy',
	    'main', 'source')), 'Sources Sources.gz')
	self.assertEqual(open(os.path.join(self.archivedir, 'dists', 'edgy', 'main',
	    'binary-i386', 'Packages')).read(), '')
	map = debcontrol_map(os.path.join(self.archivedir, 'dists', 'edgy', 'main',
	    'source', 'Sources'))

	# check generated Sources validity
	self.assertEqual(_mapkeystr(map), 'bar foo')
	self.assertEqual(map['bar']['Version'], '2-2')
	self.assertEqual(map['bar']['Directory'], 'pool/main/b/bar')
	self.assert_(map['bar']['Files'].find('bar_2-2.tar.gz') > 0)
	self.assert_(map['bar']['Files'].find('bar_2-2.dsc') > 0)
	self.assertEqual(map['foo']['Version'], '1-1')
	self.assertEqual(map['foo']['Directory'], 'pool/main/f/foo')
	self.assert_(map['foo']['Files'].find('foo_1.orig.tar.gz') > 0)
	self.assert_(map['foo']['Files'].find('foo_1-1.diff.gz') > 0)
	self.assert_(map['foo']['Files'].find('foo_1-1.dsc') > 0)

    def test_create_indexes_altext(self):
	'''create_indexes() with a non-default file name extension.'''

	# install some test debs
	self.assert_(install_pool(self.archivedir, self.create_deb('foo1',
	    '1-1', 'i386', self.testfiles_dir, '.ddeb'), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('foo2',
	    '1-1', 'i386', self.testfiles_dir, '.ddeb'), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar1',
	    '2-2ubuntu1', 'powerpc', self.testfiles_dir, '.ddeb'), 'main', 'bar'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar2',
	    '2-2ubuntu1', 'powerpc', self.testfiles_dir, '.ddeb'), 'main', 'bar'))

	# call create_indexes()
	create_indexes(self.archivedir, 'edgy', 'main restricted universe', 
	    'i386 powerpc amd64', ext='.ddeb', stderr=subprocess.PIPE)

	# check directory structure
	self.assertEqual(_ls(self.archivedir), 'db dists pool')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy',
	    'main', 'binary-i386')), 'Packages Packages.gz')
	self.assertEqual(_ls(os.path.join(self.archivedir, 'dists', 'edgy',
	    'universe', 'binary-powerpc')), 'Packages Packages.gz')

	# check Packages contents
	map = debcontrol_map(os.path.join(self.archivedir,
	    'dists', 'edgy', 'universe', 'binary-i386', 'Packages.gz'))
	self.assertEqual(map['foo1']['Architecture'], 'i386')
	self.assertEqual(map['foo1']['Version'], '1-1')
	self.assertEqual(map['foo1']['Filename'], 'pool/universe/f/foo/foo1_1-1_i386.ddeb')
	self.assert_(not map.has_key('bar1')) # since that's powerpc

	self.assertEqual(open(os.path.join(self.archivedir, 'dists', 'edgy', 'main',
	    'binary-i386', 'Packages')).read(), '')

	self.assert_(os.path.getsize(os.path.join(self.archivedir, 'dists',
            'edgy', 'Release')) > 0)

    def test_create_indexes_performance(self):
	'''performance of create_indexes()'''

	# install some test debs
        for i in xrange(1000):
            install_pool(self.archivedir, self.create_deb(str(i), '1-1',
                'i386', self.testfiles_dir), 'main', str(i))

        print '[setup done]',
        sys.stdout.flush()

	# call create_indexes()
        t1 = time.time()
	create_indexes(self.archivedir, 'edgy', 'main', 'i386', stderr=subprocess.PIPE)
        t2 = time.time()

        pkg = os.path.join(self.archivedir, 'dists', 'edgy', 'main', 'binary-i386', 'Packages')
        print 'Packages:', os.path.getsize(pkg), 'Packages.gz:', os.path.getsize(pkg+'.gz'),
        print '%.1f seconds ' % (t2-t1),
        sys.stdout.flush()

    def test_create_indexes_refarchive(self):
	'''archive_cleanup() with a reference archive'''

	# install some test debs

        # ... crufty old version
	self.assert_(install_pool(self.archivedir, self.create_deb('foo-dbgsym',
	    '1-0', 'i386', self.testfiles_dir, '.ddeb'), 'universe', 'foo'))

        # ... dapper
	self.assert_(install_pool(self.archivedir, self.create_deb('foo-dbgsym',
	    '1-1', 'i386', self.testfiles_dir, '.ddeb'), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar1-dbgsym',
	    '2-2', 'powerpc', self.testfiles_dir, '.ddeb'), 'main', 'bar'))

        # ... edgy
	self.assert_(install_pool(self.archivedir, self.create_deb('foo-dbgsym',
	    '1-2', 'i386', self.testfiles_dir, '.ddeb'), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar1-dbgsym',
	    '2-2ubuntu1', 'powerpc', self.testfiles_dir, '.ddeb'), 'main', 'bar'))

        ref_archive = tempfile.mkdtemp()
        try:
            # create reference indexes
            indexes = {
                'dapper': { 
                    'main':     { 'powerpc': 'Package: bar1\nSource:bar\nVersion: 2-2\n' },
                    'universe': { 'i386': 'Package: foo\nVersion: 1-1\n' }
                },
                'edgy': { 
                    'main':     { 'powerpc': 'Package: bar1\nSource:bar\nVersion: 2-2ubuntu1\n' },
                    'universe': { 'i386': 'Package: foo\nVersion: 1-2\n' }
                }
            }
            for release in ('dapper', 'edgy'):
                for comp in ('main', 'universe'):
                    for arch in ('i386', 'powerpc'):
                        dir = os.path.join(ref_archive, 'dists', release,
                            comp, 'binary-' + arch)
                        os.makedirs(dir)
                        p = os.path.join(dir, 'Packages')
                        f = open(p, 'w')
                        try:
                            f.write(indexes[release][comp][arch])
                        except KeyError:
                            pass
                        f.close()
                        _pack(p, p + '.gz')

            # create our indexes
            create_indexes(self.archivedir, 'dapper', 'main universe', 
                'i386 powerpc', stderr=subprocess.PIPE, ext='.ddeb',
                ref_archive='file://' + ref_archive)
            create_indexes(self.archivedir, 'edgy', 'main universe', 
                'i386 powerpc', stderr=subprocess.PIPE, ext='.ddeb',
                ref_archive='file://' + ref_archive)

            # check correct versions in Packages files

            map = debcontrol_map(os.path.join(self.archivedir,
                'dists', 'dapper', 'universe', 'binary-i386', 'Packages.gz'), 'Version')
            self.assertEqual(map.keys(), ['1-1'])
            self.assertEqual(map['1-1']['Package'], 'foo-dbgsym')

            map = debcontrol_map(os.path.join(self.archivedir,
                'dists', 'dapper', 'main', 'binary-powerpc', 'Packages.gz'), 'Version')
            self.assertEqual(map.keys(), ['2-2'])
            self.assertEqual(map['2-2']['Package'], 'bar1-dbgsym')

            map = debcontrol_map(os.path.join(self.archivedir,
                'dists', 'edgy', 'universe', 'binary-i386', 'Packages.gz'), 'Version')
            self.assertEqual(map.keys(), ['1-2'])
            self.assertEqual(map['1-2']['Package'], 'foo-dbgsym')

            map = debcontrol_map(os.path.join(self.archivedir,
                'dists', 'edgy', 'main', 'binary-powerpc', 'Packages.gz'), 'Version')
            self.assertEqual(map.keys(), ['2-2ubuntu1'])
            self.assertEqual(map['2-2ubuntu1']['Package'], 'bar1-dbgsym')

        finally:
            shutil.rmtree(ref_archive)

    def test_debcontrol_dominate(self):
	'''debcontrol_dominate()'''

	i = tempfile.NamedTemporaryFile()
	i.write('''Package: foo
Version: 1-0
Description: oldest foo

Package: bar
Version: 1-0
Description: barbar

Package: foo
Version: 1:1-10
Description: newest foo
 I wanna have this!

Package: foo
Version: 1:1-5
Description: slightly older foo 
 (needs proper VersionCompare)

Package: foo
Version: 3-1
Description: old foo crap
''')
	i.flush()

	o = StringIO()
	debcontrol_dominate(open(i.name), o)
    
	self.assertEqual(o.getvalue(), '''Package: bar
Version: 1-0
Description: barbar

Package: foo
Version: 1:1-10
Description: newest foo
 I wanna have this!
''')

    def test_archive_cleanup(self):
	'''archive_cleanup()'''

	# install some test debs
	self.assert_(install_pool(self.archivedir, self.create_deb('foo1',
	    '1-1', 'i386', self.testfiles_dir), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('foo2',
	    '1-1', 'i386', self.testfiles_dir), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar1',
	    '2-2ubuntu1', 'powerpc', self.testfiles_dir), 'main', 'bar'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar2',
	    '2-2ubuntu1', 'powerpc', self.testfiles_dir), 'main', 'bar'))

	# call create_indexes()
	create_indexes(self.archivedir, 'edgy', 'main restricted universe', 
	    'i386 powerpc amd64', stderr=subprocess.PIPE)

	# make file system snapshot, clean up, check that nothing was removed
	clean_file_list = _find_files(self.archivedir)
	archive_cleanup(self.archivedir)
	self.assertEqual(_find_files(self.archivedir), clean_file_list,
	    'missing: ' + str(clean_file_list.difference(_find_files(self.archivedir))))

	# now add some cruft
	self.assert_(install_pool(self.archivedir, self.create_deb('foo1',
	    '1-0', 'i386', self.testfiles_dir), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar1',
	    '2-2', 'powerpc', self.testfiles_dir), 'main', 'bar'))
	self.assert_(install_pool(self.archivedir, self.create_deb('cruft',
	    '5', 'amd64', self.testfiles_dir), 'main', 'cruft'))

        # test with keeping old files around; make cruft a recent and bar a > 2
        # days old deb
        bar_path = os.path.join(self.archivedir, 'pool', 'main', 'b', 'bar',
            'bar1_2-2_powerpc.deb')
        os.utime(bar_path, (time.time(), time.time()-2.1*24*3600))

	archive_cleanup(self.archivedir, keep_days=2)
        self.assert_(os.path.exists(os.path.join(self.archivedir, 'pool',
            'main', 'c', 'cruft', 'cruft_5_amd64.deb')))
        self.assert_(not os.path.exists(bar_path))

        # now clean up everything
	archive_cleanup(self.archivedir)
	self.assertEqual(_find_files(self.archivedir), clean_file_list)

    def test_archive_cleanup_customlayout(self):
	'''archive_cleanup() with a non-standard archive structure'''

	open(os.path.join(self.archivedir, 'foo_1_i386.deb'), 'w').write('foo_1')
	open(os.path.join(self.archivedir, 'foo_2_i386.deb'), 'w').write('foo_2')
	open(os.path.join(self.archivedir, 'bar_1_i386.deb'), 'w').write('bar_1')
	open(os.path.join(self.archivedir, 'cruft_3_i386.deb'), 'w').write('cruft')
	open(os.path.join(self.archivedir, 'Packages'), 'w').write('''Package: foo
Version: 2
Filename: foo_2_i386.deb

Package: bar
Version: 1
Filename: ./bar_1_i386.deb
''')
	archive_cleanup(self.archivedir)
	self.assertEqual(_find_files(self.archivedir), set([
	    os.path.join(self.archivedir, 'foo_2_i386.deb'),
	    os.path.join(self.archivedir, 'bar_1_i386.deb'),
	    os.path.join(self.archivedir, 'Packages')]))

    def test_archive_dominate(self):
	'''archive_dominate()'''

	# install some test debs and call these 'dapper'
	self.assert_(install_pool(self.archivedir, self.create_deb('foo',
	    '1-0', 'i386', self.testfiles_dir), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('foo',
	    '1-1', 'i386', self.testfiles_dir), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar',
	    '2-2', 'i386', self.testfiles_dir), 'main', 'bar'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar',
	    '2-2ubuntu1', 'i386', self.testfiles_dir), 'main', 'bar'))
	self.create_source('foo', '1-0', os.path.join(self.archivedir, 'pool',
	    'universe', 'f', 'foo'))
	self.create_source('foo', '1-1', os.path.join(self.archivedir, 'pool',
	    'universe', 'f', 'foo'))
	self.create_source('bar', '2-2', os.path.join(self.archivedir,
	    'pool', 'main', 'b', 'bar'))
	self.create_source('bar', '2-2ubuntu1', os.path.join(self.archivedir,
	    'pool', 'main', 'b', 'bar'), native=True)

	create_indexes(self.archivedir, 'dapper', 'main universe', 
	    'source i386', stderr=subprocess.PIPE)

	# install some newer versions, call these 'edgy'
	self.assert_(install_pool(self.archivedir, self.create_deb('foo',
	    '1-2', 'i386', self.testfiles_dir), 'universe', 'foo'))
	self.assert_(install_pool(self.archivedir, self.create_deb('bar',
	    '3-0ubuntu1', 'i386', self.testfiles_dir), 'main', 'bar'))
	self.create_source('foo', '1-2', os.path.join(self.archivedir, 'pool',
	    'universe', 'f', 'foo'))
	self.create_source('bar', '3-0ubuntu1', os.path.join(self.archivedir,
	    'pool', 'main', 'b', 'bar'), native=True)

	create_indexes(self.archivedir, 'edgy', 'main universe', 
	    'source i386', stderr=subprocess.PIPE)

	archive_dominate(self.archivedir)

	# check that we only have the latest stuff from dapper and edgy
	self.assertEqual(_find_files(os.path.join(self.archivedir, 'pool'),
	    strip_prefix=self.archivedir), set([
	    '/pool/main/b/bar/bar_2-2ubuntu1.dsc',
	    '/pool/main/b/bar/bar_2-2ubuntu1.tar.gz',
	    '/pool/main/b/bar/bar_2-2ubuntu1_i386.deb',
	    '/pool/main/b/bar/bar_3-0ubuntu1.dsc',
	    '/pool/main/b/bar/bar_3-0ubuntu1.tar.gz',
	    '/pool/main/b/bar/bar_3-0ubuntu1_i386.deb',
	    '/pool/universe/f/foo/foo_1-1.diff.gz',
	    '/pool/universe/f/foo/foo_1-1.dsc',
	    '/pool/universe/f/foo/foo_1-1_i386.deb',
	    '/pool/universe/f/foo/foo_1-2.diff.gz',
	    '/pool/universe/f/foo/foo_1-2.dsc',
	    '/pool/universe/f/foo/foo_1-2_i386.deb',
	    '/pool/universe/f/foo/foo_1.orig.tar.gz'
	    ]))

	# check that Packages.gz and Sources.gz are re-created properly
	s = os.path.join(self.archivedir, 'dists', 'dapper', 'main', 
	    'source', 'Sources')
	self.assertEqual(open(s).read(), gzip.open(s+'.gz').read())
	s = os.path.join(self.archivedir, 'dists', 'dapper', 'main',
	    'binary-i386', 'Packages')
	self.assertEqual(open(s).read(), gzip.open(s+'.gz').read())
	s = os.path.join(self.archivedir, 'dists', 'edgy', 'universe',
	    'source', 'Sources')
	self.assertEqual(open(s).read(), gzip.open(s+'.gz').read())
	s = os.path.join(self.archivedir, 'dists', 'edgy', 'universe',
	    'binary-i386', 'Packages')
	self.assertEqual(open(s).read(), gzip.open(s+'.gz').read())

    def test_archive_dominate_flat(self):
	'''archive_dominate() with a flat tree'''

	# install some test debs and call these 'dapper'
	self.assert_(self.create_deb('foo', '1-0', 'i386', self.archivedir))
	self.assert_(self.create_deb('foo', '1-1', 'i386', self.archivedir))
	self.create_source('foo', '1-0', self.archivedir)
	self.create_source('foo', '1-1', self.archivedir)

	# create indices
	old_cwd = os.getcwd()
	try:
	    os.chdir(self.archivedir)
	    assert subprocess.call('apt-ftparchive packages . > Packages',
		shell=True) == 0
	    assert subprocess.call('apt-ftparchive sources . > Sources',
		shell=True, stderr=subprocess.PIPE) == 0
	finally:
	    os.chdir(old_cwd)

	archive_dominate(self.archivedir)

	self.assertEqual(_find_files(self.archivedir, strip_prefix=self.archivedir), set([
	    '/Packages', '/Sources', '/foo_1-1_i386.deb', '/foo_1.orig.tar.gz',
	    '/foo_1-1.dsc', '/foo_1-1.diff.gz']))

if __name__ == '__main__':
    unittest.main()