#!/usr/bin/env python

import commands
import os
import os.path
import sys

# Location of the DPKG files
dpkg_dir = '/var/lib/dpkg'

# Global variables
name_lookup = {}
rootfs_id = ''

# Read in the full display names for packages
def read_names():
    global dpkg_dir
    global name_lookup
    pkg = ''

    # Read though the DPKG status file
    pkg_stat = open(dpkg_dir + '/status', 'r')

    for line in pkg_stat:
        line = line.rstrip()

        # Extract the package name and display name
        if line[0:8] == 'Package:':
            pkg = line[9:]
        elif line[0:19] == 'Maemo-Display-Name:':
            # Only store the display name if it's different
            full_name = line[20:]

            if pkg != full_name:
                name_lookup[pkg] = full_name

    pkg_stat.close()

# Lookup the full display name for a package
def get_name(pkg):
    global name_lookup

    # Check whether a full display name exists for this package
    if name_lookup.has_key(pkg):
        return name_lookup[pkg]

    return

# Return the filesystem ID for the specified file/directory
def get_fsid(file):
    return os.stat(file).st_dev

# Check whether the specified file is on the rootfs or not
def is_rootfs(file):
    global rootfs_id
    global rootfs_pkg

    # Store the rootfs ID if it hasn't yet been set
    if rootfs_id == '':
        rootfs_id = get_fsid('/')

    # Check whether the filesystem ID matches that of the rootfs
    if get_fsid(file) == rootfs_id:
        return 1

    return 0

# Find the on-disk usage of a file
def disk_usage(file, seen):
    # Run stat on the file
    stat = os.stat(file)

    inode = stat.st_ino
    if seen.has_key(inode):
        return

    seen[inode] = 1

    blksize = stat.st_blksize       # Block size for the disk
    chunknum = stat.st_blocks       # Number of 512b blocks used by file
    blk_mult = blksize / 512        # How many 512b chunks are in one block

    # Work out how many blocks the file is actually using
    blk_count = (chunknum + blk_mult - 1) / blk_mult

    # Return the size in kb
    return blk_count * blksize / 1024

# Lookup the sizes for a given package
def get_size(pkg, list = 0):
    global dpkg_dir

    # Open the packages .list file
    try:
        pkg_info = open(dpkg_dir + '/info/' + pkg + '.list', 'r')
        dulist = {}
        seen = {}

        # Read through and store the paths for files on the rootfs
        for file in pkg_info:
            file = file.rstrip()

            if not os.path.isfile(file):
                continue
            if os.path.islink(file):
                continue
            if not is_rootfs(file):
                continue

            # Lookup the disk usage of the file
            size = disk_usage(file, seen)

            if size is None:
                continue
            if size == 0:
                continue

            dulist[file] = size
    except IOError:
        return
    finally:
        pkg_info.close()

    # Check whether there are any files on the rootfs
    if len(dulist) == 0:
        return

    # Return either the total size or the list of individual file sizes
    if list == 0:
        return sum(dulist.values())
    else:
        return dulist

# Get the total rootfs size for all non-l10n packages
def checkpkgs():
    # Store the display names for later lookup
    read_names()

    # Loop through all files in the DPKG info directory
    for file in os.listdir(dpkg_dir + '/info'):
        # Skip any that aren't .list files
        if file[-5:] != '.list':
            continue

        # Skip the l10n packages
        if file.find('l10n') > -1:
            continue

        # Extract the package name
        pkg = file[0:len(file) - 5]

        # Lookup the package rootfs size
        size = get_size(pkg)

        # Skip further processing if there's no rootfs usage
        if size is None:
            continue

        # Lookup the package's display name
        name = get_name(pkg)

        # If there is one, append it in brackets
        if not name is None:
            pkg = pkg + ' (' + name + ')'

        # Output the package size and name
        print '%d\t%s' % (size, pkg)

# List the rootfs files in the given package (and their sizes)
def pkglist(pkg):
    # Lookup the array of files and sizes
    sizes = get_size(pkg, 1)

    # Output the info
    if not sizes is None:
        for (file, size) in sizes.items():
            print "%d\t%s" % (size, file)

# Report the total rootfs usage for the given package
def pkgsize(pkg):
    # Lookup the rootfs usage
    size = get_size(pkg)

    # Report the rootfs usage (or 0 if there is none)
    if size is None:
        print '0\t%s' % pkg
    else:
        print '%d\t%s' % (size, pkg)

# Use the first parameter as the command to run.
if sys.argv[1] == 'checkpkgs':
    checkpkgs()
elif sys.argv[1] == 'pkglist':
    pkglist(sys.argv[2])
elif sys.argv[1] == 'pkgsize':
    pkgsize(sys.argv[2])

# vim: ts=4 et