#!/usr/bin/python
import rrdtool, re, time

step=10 # how many seconds between 2 data points
keep=[(10, 3600*24*10), # how many dozens of seconds of data to keep
      (180, 20*24*42), # how many 3-minutes intervals
      ]
class base:
    def create (self, rrdname):
        def _make_ds (dsn):

            # default type is COUNTER (valid for net, disks, and most cpu info)
            dst="COUNTER"

            # one disk data is of type GAUGE
            if dsn=="io_inprogress": dst="GAUGE"

            # all memory stats are of type GAUGE,
            # and they are the only one to start with an uppercase
            if dsn[0]==dsn[0].upper(): dst="GAUGE"

            # all cpu type is of type COUNTER,
            # except procs_[running|blocked]
            if dsn.startswith("procs_"): dst="GAUGE"
            
            return "DS:%s:%s:666:0:4000000000"%(dsn,dst)
        ds = [ _make_ds (dsn) for dsn in self.descs ]
        rras = [ "RRA:LAST:0:1:10" ]
        for cf in ["MIN","MAX","AVERAGE"]:
            for stepsize, stepquantity in keep:
                rras += [ "RRA:%s:0.9:%d:%d"%(cf,stepsize/step,stepquantity/step) ]
        rrdtool.create (rrdname, "-s %d"%step, *(ds+rras))
        
    def update (self, rrdname, values):
        try: open(rrdname).close()
        except: self.create (rrdname)
        template = ":".join(self.descs)
        values = ":".join([a.strip() for a in values])
        rrdtool.update (rrdname, "-t"+template, "N:"+values)

class disk(base):
    descs = ["junk", "major", "minor", "name",
             "read_issued", "read_merged", "read_sectors", "read_time",
             "write_completed", "write_merged", "write_sectors", "write_time",
             "io_inprogress", "io_time", "io_weighted" ][4:]
    def loop (self):
        f = open ("/proc/diskstats")
        for disk in f.readlines():
            values = re.split ("[ \t]+", disk)
            if len(values)==8: continue
            if values[3].startswith("ram"): continue
            self.update ("disk-%s.rrd"%values[3], values[4:])
        f.close ()

class net(base):
    descs = [ "junk", "name",
              "rx_bytes", "rx_packets", "rx_errs", "rx_drop", "rx_fifo",
              "rx_frame", "rx_compressed", "rx_multicast",
              "tx_bytes", "tx_packets", "tx_errs", "tx_drop", "tx_fifo",
              "tx_colls", "tx_carrier", "tx_compressed" ][2:]
    def loop (self):
        f = open ("/proc/net/dev")
        for interface in f.readlines()[2:]:
            values = re.split ("[ \t:]+", interface)
            self.update ("net-%s.rrd"%values[1], values[2:])
        f.close ()

class mem(base):
    descs = [ "MemTotal", "MemFree", "Buffers", "Cached", "SwapCached",
              "Active", "Inactive",
              "HighTotal", "HighFree", "LowTotal", "LowFree",
              "SwapTotal", "SwapFree", "Dirty", "Writeback",
              "Mapped", "Slab", "Committed_AS", "PageTables",
              "VmallocTotal", "VmallocUsed", "VmallocChunk" ]
    def loop (self):
        h = {}
        f = open ("/proc/meminfo")
        for item in f.readlines():
            m = re.match ("^(.*):[ \t]+([0-9]+) kB$", item)
            if not m:
                print "Couldn't parse /proc/meminfo - %s"%repr(item)
                continue
            desc, value = m.groups()
            h[desc] = str(1024*int(value))
        self.update ("meminfo.rrd", [h[k] for k in self.descs])
        f.close ()

class cpu(base):
    basedescs = ["ctxt", "processes", "procs_running", "procs_blocked"]
    cpudescs = ["user", "nice", "system", "idle"]
    def __init__ (self, howmanycpus=1):
        self.descs = self.basedescs[:]
        for cpu in [""]+[str(x) for x in range(howmanycpus)]:
            self.descs += ["cpu%s_%s"%(cpu, subdesc) for subdesc in self.cpudescs]
    def loop (self):
        h = {}
        f = open ("/proc/stat")
        for item in f.readlines():
            items = re.split ("[ \t]+", item)
            key, values = items[0], items[1:]
            if key in self.basedescs:
                h[key] = values[0]
                continue
            if key.startswith("cpu"):
                for i,subdesc in zip(range(len(self.cpudescs)),self.cpudescs):
                    h["%s_%s"%(key,subdesc)] = values[i]
                continue
        self.update ("cpu.rrd", [h[k] for k in self.descs])
        f.close ()
        
def mainloop ():
    mynet = net ()
    mydisk = disk ()
    mymem = mem ()
    mycpu = cpu ()
    while 1:
        mynet.loop()
        mydisk.loop()
        mymem.loop()
        mycpu.loop()
        time.sleep(step)

mainloop()