#!/bin/env python

def docs():
   print """
duscan.py
Filesystem scanning utility
--------------------------
duscan.py is designed to scan a filesystem to catalogue where the memory is
being used. In each directory it write a small .du file with that directory's
information. In the top directory, it also writes .du_file, .du_dir, .du_user,
and .du_group, which contain information about the whole directory tree.

Running duscan
The basic duscan command is python duscan.py. It will run on python 2.4 through
2.6, but has not been tested on earlier or later versions. This will generate a
.du file in  every directory and .du_file/dir/user/group in the top directory. 

duscan takes the following options: abdDhiopPtuv

-a (--add) will scan the tree normally, but if a .du file is found that is less
    than a certain age, it will use the information from that file rather than
    scanning that subtree. It will attempt to read the .du_* files from  the
    same directory, but if it can't, the final .du_* files will be incomplete.

-b (--debug) run in debug mode. Prints out a ton of information, most of it not
    very useful

-d (--dir) set the top directory (where the scan will start from). Defaults to
    the current directory if not set.

-h (--help) print these docs and exit

-i (--img) generate .png pie graphs of the directories and put one file in each
    directory in the tree. Slows the program down a lot. If using this option,
    duscan should be run through sage or ipython, or another version of python
    that includes the pylab libraries

-o (--onlyvisible) only follow visible directories (i.e. not ones starting with
    a ".")

-p (--html-prefix) for use generating html. This is the prefix that will be
    built into the html links

-P (--fs-prefix) the part of the file system not to be included in links
    generated  by the html

-t (--html) generate html files (.du.html) with the duscan information.

-u (--use) use the image generated by -i in the html

-v (--verbose) run in verbose mode. Less information than debug mode, but prints
    a message when entering and exiting directories, as well as when writing a
    .du file and when reading an existing .du file (if running with -a)

duscan should be run as root, otherwise it will exit with an error message if
permission is denied to a directory.

EXCLUDE: list of the absolute paths to directories that you do not want scanned
(no trailing slash)

SUPPRESS: list of the absolute paths to directories you want scanned, but not
written to (no trailing slash)

Files:
******************
.du
Format:

TS FS DS TE FE DE SF SD TF TD
D1 S1
D2 S2
...

with the following meanings:
TS - Total Size of the directory (FS + DS) (MB)
FS - Total size of files in the directory (MB)
DS - Total size of the subdirectories (MB)
TE - Total number of entries in the directory
FE - Total number of file entries
DE - Total number of directory entries
SF - Total number of symlinked files
SD - Total number of symlinked directories
TF - Accumulated number of files (including all subdirs)
TD - Accumulated number of directories
D1/S1 - The largest subdirectory/its size (MB)
D2/S2 - Second largest subdirectory/its size (MB)
... - the rest of the directories sorted by size 
******************
.du_file/.du_dir
These files list all files greater than 10 MB and all directories larger than
100 MB (though these constants can easily be changed). 
Format:

SIZE1 ABSPATH1
SIZE2 ABSPATH2

listed in no particular order.
******************
.du_user/.du_group
These files list the total space used by each user/group on the system.

Format:
USERNAME/GROUPNAME SIZE

also in no particular order
******************
"""

import os
import os.path
import sys

import stat
import logging
import pwd
import grp
import time

from optparse import OptionParser

EXCLUDE         = [ "/proc",
                    "/sys",
                    "/dev",
                    "/Volumes",
                    "/etc",
                    "/var/lib/nfs",
                    ]

SUPPRESS        = []

SCALE           = 2.0**20   # use MB -- stat outputs size in bytes
#minimum size to be put in list of big files/directories
MIN_FILE_SIZE   = 10        # MB       
MIN_DIR_SIZE    = 100       # MB

PIE_MIN_PCT     = 0.02      # less than PIE_MIN_PCT, don't show details

MAX_REC_DEPTH   = 50        # safety to stop recursion after this depth
MIN_TIME        = 1         # seconds since .du last changed 24*60*60 #24 hours
MIN_ADD_TIME    = 24*60*60  # .du file must be less than a day old to add to it
SLEEP_TIME      = 0.01      # seconds to sleep for each iteration (safety)

PROC_UID        = os.getuid()

dirstack=[]

parser = OptionParser()
parser.add_option("-d", "--dir", dest="topdir", default=os.getcwd(),
                  help="Set the starting directory.", metavar="DIR")
parser.add_option("-t", "--html", dest="gen_html", action="store_true",
                  help="Generate .du.html files in each directory", default=False)
parser.add_option("-i", "--img",dest="gen_img", action="store_true",
                  help="Generate pie graphs for each directory.", default=False)
parser.add_option("-u", "--use",dest="use_img", action="store_true",
                  help="Include pie graphs in html.", default=False)
parser.add_option("-v", "--verbose",dest="verbose", action="store_true",
                  help="Run program in verbose mode.", default=False)
parser.add_option("-o", "--onlyvisible",dest="only_vis", action="store_true",
                  help="Only follow paths through visible directories.", default=False)
parser.add_option("-b", "--debug",dest="debug", action="store_true",
                  help="Set to debug mode.", default=False)
parser.add_option("-p","--html-prefix",dest="html_prefix",default="", help="Url prefix to be built into html.",
                  metavar="PREFIX")
parser.add_option("-P","--fs-prefix",dest="fs_prefix", help="File system prefix.",
                  default="", metavar="PREFIX")
parser.add_option("-D","--docs",dest="show_docs",help="Show documentation", default=False,
                  action="store_true")
parser.add_option("-a","--add",dest="add",help="Add to a previous scan.", default=False,
                  action="store_true")

(options,args) = parser.parse_args()

if options.debug:
    logging.basicConfig(level=logging.DEBUG)
elif options.verbose:
    logging.basicConfig(level=logging.INFO)
else:
    logging.basicConfig(level=logging.WARNING)

if options.gen_img:
    import pylab as pl

TOPDIR = options.topdir
logging.debug("TOPDIR " + TOPDIR)


if options.show_docs:
    docs()
    sys.exit(0)

def fix_perms(fn):
    """ fn needs to be a bare file name, no directories.
        It must be in the current directory """
    os.chmod(fn, 0644) # this is for everyone
    if PROC_UID == 0:
        # we're root, so chown the .du* files to the owner of the current directory
        fn_stat     = os.stat(os.getcwd())
        uid, gid    = fn_stat[stat.ST_UID], fn_stat[stat.ST_GID]
        os.chown(fn, uid, gid)

def writedu(stats):
    """Writes the .du file"""
    fh = open(".du", "w")
    size_stats = "%.2f %.2f %.2f %d %d %d %d %d %d %d\n" %(stats["total_size"], stats["file_size"], stats["dir_size"],
                                                            stats["total"], stats["rfiles"], stats["rdirs"], stats["lfiles"], 
                                                            stats["ldirs"], stats["total_files"], stats["total_dirs"])
    fh.write(size_stats)
    stats["dirs"].sort(reverse=True)
    for dir in stats["dirs"]:
        fh.write("%.2f\t%s\n" %(dir[0], dir[1]))
    fh.close()
    fix_perms(fh.name)
 
def generateIMG(stats):
    """Generates a pie chart of percentages of space used based on .du data"""
    labels      = []
    fracs       = []
    sizeleft    = stats["total_size"]
    try:
        for dir in stats["dirs"]:
            if dir[0]/stats["total_size"] >= PIE_MIN_PCT:
                labels.append(dir[1])
                fracs.append(dir[0]/stats["total_size"])
                sizeleft -= dir[0]
        if stats["file_size"]/stats["total_size"] >= PIE_MIN_PCT:
            labels.append("files")
            fracs.append(stats["file_size"]/stats["total_size"])
            sizeleft -= stats["file_size"]
        if sizeleft/stats["total_size"] >= PIE_MIN_PCT:
            labels.append("other")
            fracs.append(sizeleft/stats["total_size"])
    except ZeroDivisionError:
        pass
    
    pl.rcParams['font.size']=7.0
    pl.figure(figsize=(5,4))
    pl.ax = pl.axes([0.2, 0.2, 0.56, 0.7])
    pl.pie(fracs, labels=labels, autopct='%1.1f%%', shadow=True)
    pl.title("Storage")
    pl.savefig(".du.png")
    pl.close('all')
    fix_perms(".du.png")
    
    #logging.debug(fracs)
    
def writeHtmlTitle(dp,html_fh):
    rel_dp = os.getcwd()[len(options.fs_prefix):]
    dirs = rel_dp.split("/")
    link_dp = ""
    part_dp = ""
    for dir in dirs:
        #logging.debug("DIR: +%s+" %dir)
        if dir:
            link_dp += "<a href=%s/%s/.du.html>%s/</a>" %(options.html_prefix,dir,dir)
    
    #logging.debug(link_dp)
    html_fh.write("""<h2>Stats for directory %s/%s</h2>""" %(options.fs_prefix,link_dp))
    
def writehtml(dp,stats):
    """Writes the html file for the .du output"""
    html_fh = open(".du.html","w")
    html_fh.write("<a href='.'>Listing</a> | <a href='.du'>.du</a>\n ")
    rel_dp = TOPDIR.replace(options.fs_prefix,options.html_prefix,1)
    if not dp==TOPDIR:
        html_fh.write(""" | <a href="../.du.html">Up one level</a> | <a href="%s/.du.html">Top</a>""" %(rel_dp))
    
    if options.use_img:
        html_fh.write("""
<table>
    <tr>
        <td valign="top">
""")
        
    html_fh.write("""
  <table>
    <tr><td>Total size:</td><td>%.2f</td></tr>
    <tr><td>File size:</td><td>%.2f</td></tr>
    <tr><td>Directory size:</td><td>%.2f</td></tr>
    <tr><td>Total number of entries:</td><td>%d</td></tr>
    <tr><td>Number of files:</td><td>%d</td></tr>
    <tr><td>Number of directories:</td><td>%d</td></tr>
    <tr><td>Number of symlinked files:</td><td>%d</td></tr>
    <tr><td>Number of symlinked directories:</td><td>%d</td></tr>
    <tr><td>Number of files in all subdirs:</td><td>%d</td></tr>
    <tr><td>Number of dirs in all subdirs:</td><td>%d</td></tr>
  </table>
  <br>
""" %(stats["total_size"],stats["file_size"],stats["dir_size"],
      stats["total"], stats["rfiles"], stats["rdirs"], stats["lfiles"], 
      stats["ldirs"],stats["total_files"],stats["total_dirs"]))
    
    if options.use_img:
        html_fh.write("""
        </td>
        <td><img src=".du.png"></td>
    </tr>
</table>
""")
    
    html_fh.write("<br><br>")
    
    if stats["dirs"]:
        html_fh.write("""
  <table cellpadding="5" >
    <thead><tr><th align="left">dir</th><th>listing</th><th>size (MB)</th></tr></thead>
    <tbody>
""")
        for dir in stats["dirs"]:
            html_fh.write("""
<tr>
    <td><a href="%s/.du.html">%s</a></td>
    <td><a href="%s">listing</a></td>
    <td>%.2f</td>
</tr>
""" %(dir[1],dir[1],dir[1],dir[0]))
        html_fh.write("""  </tbody></table>""")
    html_fh.close()
    fix_perms(html_fh.name)
        
def genInitHtml(dp, contents):
    logging.debug("Initial html")
    
    dirs = []
    for c in contents:
        fp = "%s" %(c) 
        if os.path.isdir(fp):
            #logging.debug("dirname: %s" %c)
            dirs.append(c)
    
    html_fh = open(".du.html", "w")
    writeHtmlTitle(dp,html_fh)
    #logging.debug(dirs)
    html_fh.write("<h3>sub directories</h3>\n<p>")
    if not dirs:
        html_fh.write("""None""")
    for dir in dirs:
        html_fh.write("<a href='%s/.du.html'>%s</a>\n" %(dir,dir))
    html_fh.write("</p>")
    html_fh.close()
    fix_perms(html_fh.name)

def recursion_check(contents,dp,rec_depth):
    """Some verification methods to prevent infinite recursion: checks recursion depth and
    modification time. Also returns False for an empty directory. Adds sleep time to make debugging
    easier."""
    
    #time.sleep(SLEEP_TIME) # This was a defensive measure to stop runaway process.
    #logging.debug("Recursion depth: %d" %rec_depth)
    if rec_depth > MAX_REC_DEPTH:
        logging.error("Maximum recursion depth [%d] exceeded in [%s]" % (rec_depth, dp))
        return False
    if not contents:
        return False
    du_fp = "%s/.du" %dp
    if os.path.exists(du_fp):
        cur_time = time.time()
        du_stats = os.lstat(du_fp)
        if cur_time - du_stats[stat.ST_MTIME] < MIN_TIME:
            logging.error("File too recently modified in directory %s" %dp)
            logging.error("Now: %d Last Modified: %d" %(cur_time,du_stats[stat.ST_MTIME]))
            return False
    return True

def read_du(file_fh, dir_fh, users, groups):
    try:
        du_fh = open(".du")
        line  = du_fh.readline()[:-1].split()
        stats = (float(line[0]),float(line[8]),float(line[9]))
    except:
        return None

    for fn in [".du_file", ".du_dir"]:
        if os.path.exists(fn):
            try:
                fh = open(fn)
                file_fh.write(fh.read())
                fh.close()
            except:
                pass

    for dict in [users,groups]:
        if dict == users:
            fn = ".du_user"
        else:
            fn = ".du_group"
        if os.path.exists(fn):
            try:
                fh    = open(fn)
                lines = fh.readlines()
                for line in lines:
                    parts = line[:-1].split()
                    if not dict.has_key(parts[0]):
                        dict[parts[0]] = 0.0
                    dict[parts[0]] += float(parts[1])
            except:
                pass
        return stats

def scan(file_fh, dir_fh, users, groups, rec_depth, inodes, suppress=False):
    stats = {"total_size" : 0, "file_size" : 0, "dir_size" : 0, "dirs": [],
             "total" : 0, "rfiles" : 0, "rdirs" : 0, "lfiles" : 0, "ldirs" : 0,
             "total_files": 0, "total_dirs": 0}
    
    cur_dp = os.getcwd()
    logging.info("Entering directory: %s" %cur_dp)
    
    if options.add and os.path.exists(".du"):
        cur_time = time.time()
        if cur_time - os.lstat(".du")[stat.ST_MTIME] < MIN_ADD_TIME:
            stats = read_du(file_fh,dir_fh,users, groups)
            if stats != None:
                return stats

    contents = os.listdir(cur_dp)
    logging.debug(contents)
    contents.sort()
    if not recursion_check(contents,cur_dp,rec_depth):
        return (0,0,0)
    
    if cur_dp in SUPPRESS:
        suppress = True
    if not suppress:
        try:
            temp_fh = open(".du", "w")
            temp_fh.close()
        except:
            logging.error("Error writing to directory [%d]" % cur_dp)
            return (0,0,0)

    if options.gen_html:
        genInitHtml(dp, contents)
    if not contents:
        logging.debug("Empty Dir.")
        return (0,0,0)

    for fn in contents:
        abs_fp = os.path.abspath(fn)
        file_stats = {}

        try:
            file_stats = os.lstat(abs_fp)
        except:
            warning("File [%s] does not exist. It has probably been recently removed." % abs_fp)
            continue

        logging.debug("filename: %s" %fn)

        if (fn[0] == "." and options.only_vis and os.path.isdir(fn)) or abs_fp in EXCLUDE:
            if os.path.isdir(fn):
                stats["rdirs"]  += 1
                stats["total_files"] +=1
            elif os.path.isfile(fn):
                stats["rfiles"] += 1
                stats["total_dirs"] += 1
        elif os.path.isdir(fn) and not os.path.islink(fn):
            #prevents revisiting nodes
            logging.debug("considering [%s]" % abs_fp)
            if not file_stats[stat.ST_INO] in inodes:
                try:
                    inodes.add(file_stats[stat.ST_INO])
                    dirstack.append(cur_dp)
                    os.chdir(fn)
                    (dir_size,ftotal,dtotal) = scan(file_fh,dir_fh, users, groups, rec_depth+1, inodes,suppress)
                #logging.debug("dir: %s size: %f" %(fn,dir_size))
                    os.chdir(dirstack.pop())
                except:
                    logging.warn("Scan failed for [%s]" % fn)
                    continue
                
                if dir_size:
                    stats["dir_size"] += dir_size
                    stats["dirs"].append((dir_size,fn)) #("%.2f %s\n" %(dir_size,fn))
                    
                stats["rdirs"]          += 1
                stats["total_files"]    += ftotal
                stats["total_dirs"]     += dtotal + 1
                if dir_size > MIN_DIR_SIZE:
                    dir_fh.write("%.2f\t%s\n" %(dir_size,abs_fp))
        elif os.path.isfile(fn) and not os.path.islink(fn):
            file_size               = file_stats[stat.ST_SIZE]/SCALE
            stats["file_size"]      += file_size
            stats["rfiles"]         += 1
            stats["total_files"]    += 1
            if file_size > MIN_FILE_SIZE:
                file_fh.write("%.2f\t%s\n" %(file_size,abs_fp))
                
            try:
                username = pwd.getpwuid(file_stats[stat.ST_UID]).pw_name
            except KeyError:
                username = file_stats[stat.ST_UID]
            try:
                groupname = grp.getgrgid(file_stats[stat.ST_GID]).gr_name
            except KeyError:
                groupname = file_stats[stat.ST_GID]
                
            if not users.has_key(username):
                users[username] = 0.0
            users[username] += file_size
            if not groups.has_key(groupname):
                groups[groupname] = 0.0
            groups[groupname] += file_size
            
        elif os.path.isdir(fn) and os.path.islink(fn):
            stats["ldirs"]      += 1
        elif os.path.isdir(fn) and os.path.islink(fn):
            stats["lfiles"]     += 1
        
        stats["total"]          += 1
    
    stats["total_size"] = stats["file_size"] + stats["dir_size"]
    
    if not suppress:
        writedu(stats)
        logging.info("writing .du in %s" %cur_dp)
    
    if options.gen_html:
        #logging.info("writing .du.html in %s" %cur_dp)/opt/osg-shared/home/site/ccolgrove/workspace/scripts
        writehtml(dp,stats)
        
    if options.gen_img:
        #logging.info("writing .du.png in %s" %cur_dp)
        generateIMG(stats)
    
    logging.info("Exiting directory %s" %cur_dp)
    return (stats["total_size"],stats["total_files"],stats["total_dirs"])

users       = {}
groups      = {}
inodes      = set()

start_dir = os.getcwd()
os.chdir(TOPDIR)
dirstack.append(options.fs_prefix)
logging.debug("CWD: %s" %os.getcwd())
logging.debug(options.html_prefix)

fmode = "w"
dmode = "w"
if options.add:
    if os.path.exists(".du_file"):
        fmode = "a"
    if os.path.exists(".du_dir"):
        dmode = "a"

file_fh     = open(".du_file", fmode)
dir_fh      = open(".du_dir", dmode)

scan(file_fh, dir_fh, users, groups, 0, inodes)
file_fh.close()
dir_fh.close()
topdinfo = os.stat(".") # we're now in TOPDIR, or at least we should be!

fix_perms(file_fh.name)
fix_perms(dir_fh.name)



user_fh     = open(".du_user", "w")
group_fh    = open(".du_group", "w")
for user in users.keys():
    user_fh.write("%s\t%.2f\n"       % (user,users[user]))
for group in groups.keys():
    group_fh.write("%s\t%.2f\n" %(group,groups[group]))
user_fh.close()
group_fh.close()

fix_perms(user_fh.name)
fix_perms(group_fh.name)

html_fh = open(".du.html", "a")
html_fh.write("""
<p>Accumulated stats:
<a href=".du_file">Large Files</a>      |
<a href=".du_dir">Large Dirs</a>        |
<a href=".du_user">User storage</a>     |
<a href=".du_group">Group storage</a>
</p>""")
html_fh.close()
fix_perms(html_fh.name)

print os.getcwd()
os.chdir(start_dir)
