#!/usr/bin/env python
import os
import sys
import commands

username = os.getlogin ()
my_pid = os.getpid ()


# Runs a command a returns a sequence of strings, one string per line
# of output.  If the command exits unsuccessfully, returns None.
def get_lines_from_command (command):
    (status, output) = commands.getstatusoutput (command)
    if os.WIFEXITED (status) and os.WEXITSTATUS (status) == 0:
        return output.splitlines ()
    else:
        return None

# Takes a sequence of strings which are the output lines of "ps aux",
# and returns a sequence of integer pids that correspond to the
# current user's processes.
def get_user_pids (lines):
    result = []

    for l in lines:
        fields = l.split ()
        user = fields[0]

#        if user == username:
        if user != "USER":
            pid = int (fields[1])
            result.append (pid)

    return result

class Mapping:
    def __init__ (self, size, rss, shared_clean, shared_dirty, private_clean, private_dirty, permissions, name):
        self.size = size
        self.rss = rss
        self.shared_clean = shared_clean
        self.shared_dirty = shared_dirty
        self.private_clean = private_clean
        self.private_dirty = private_dirty
        self.permissions = permissions
        self.name = name
        self.count = 0

# Parses a line of the form "foo: 42 kB" and returns an integer for the "42" field
def parse_smaps_size_line (line):
    # Rss:                  8 kB
    fields = line.split ()
    return int(fields[1])

# Parses the string contents of /proc/1234/smaps and returns a list of Mapping objects
def parse_smaps (input):
    mappings = []

    lines = input.splitlines ()
    num_lines = len (lines)
    line_idx = 0

    # 08065000-08067000 rw-p 0001c000 03:01 147613     /opt/gnome/bin/evolution-2.6
    # Size:                 8 kB
    # Rss:                  8 kB
    # Shared_Clean:         0 kB
    # Shared_Dirty:         0 kB
    # Private_Clean:        8 kB
    # Private_Dirty:        0 kB

    while num_lines > 0:
        fields = lines[line_idx].split (" ", 5)
        if len (fields) == 6:
            (offsets, permissions, bin_permissions, device, inode, name) = fields
        else:
            (offsets, permissions, bin_permissions, device, inode) = fields
            name = ""

        size          = parse_smaps_size_line (lines[line_idx + 1])
        rss           = parse_smaps_size_line (lines[line_idx + 2])
        shared_clean  = parse_smaps_size_line (lines[line_idx + 3])
        shared_dirty  = parse_smaps_size_line (lines[line_idx + 4])
        private_clean = parse_smaps_size_line (lines[line_idx + 5])
        private_dirty = parse_smaps_size_line (lines[line_idx + 6])

        name = name.strip ()

        mapping = Mapping (size, rss, shared_clean, shared_dirty, private_clean, private_dirty, permissions, name)
        mappings.append (mapping)

        num_lines -= 7
        line_idx += 7

    return mappings

# Returns a list of Mapping objects for the specified pid
def get_mappings_for_pid (pid):
    try:
        smaps_file = open ("/proc/%s/smaps" % pid, "r")
    except:
        return None

    smaps = smaps_file.read ()
    smaps_file.close ()

    mappings = parse_smaps (smaps)
    return mappings

def compare_totals (a, b):
    return cmp (a[1], b[1])

def print_map (name, size, num_mappings):
    print "%s: %s KB (spread among %s mappings)" % (name, size, num_mappings)

psaux_lines = get_lines_from_command ("ps aux")
pids = get_user_pids (psaux_lines)

name_to_total = {}

for pid in pids:
    if pid == my_pid:
        continue

    mappings = get_mappings_for_pid (pid)
    if not mappings:
        continue

    for m in mappings:
        if name_to_total.has_key (m.name):
            total = name_to_total[m.name]
        else:
            total = Mapping (0, 0, 0, 0, 0, 0, 0, m.name)
            name_to_total[m.name] = total

        total.shared_dirty += m.shared_dirty
        total.private_dirty += m.private_dirty
        total.count += 1

# Private

print "PRIVATE_DIRTY:"

totals = []

for n in name_to_total:
    t = name_to_total[n]
    totals.append ((t.name, t.private_dirty, t.count))

totals.sort (compare_totals)

final = 0
libraries = 0

for t in totals:
    if t[1] != 0:
        print_map (t[0], t[1], t[2])
        final += t[1]
        if len (t[0]) > 0 and t[0][0] == '/':
            libraries += t[1]

print "Private_dirty total: %s KB" % final
print "Private_dirty for file mappings: %s KB" % libraries

# Shared

print "\nSHARED_DIRTY:"

totals = []

for n in name_to_total:
    t = name_to_total[n]
    totals.append ((t.name, t.shared_dirty, t.count))

totals.sort (compare_totals)

final = 0

for t in totals:
    if t[1] != 0:
        print_map (t[0], t[1], t[2])
        final += t[1]

print "Shared_dirty total: %s KB" % final