'''
Sagators statistics collector and distributor, version 0.3.2

(c) 2003-2004 Jan ONDREJ (SAL) <ondrejj(at)salstar.sk>
                                                                                
 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.

Examples:
  python2 sagator.py /
  python2 sagator.py mrtg Total-count Filtered-count

'''

import socket,re,string,os,sys
import urllib,time,signal,struct
import interscan.match
from avlib import *

__all__ = ['collector','update','query','status']

COLLECTOR_SERVER=()
MYCOLLECTOR=None

def uptime(starttime):
    up=time.mktime(time.localtime())-starttime
    mins=int(up/60)%60
    hours=int(up/3600)%24
    days=int(up/86400)
    if days>0:
      return "%d days, %02d:%02d" % (days,hours,mins)
    else:
      return "%02d:%02d" % (hours,mins)

class collector:
  pid=0
  data={ 'Total-count':0.0, 'Total-bytes':0.0,
         'Filtered-count':0.0, 'Filtered-bytes':0.0,
         'Total-time':0.0 }
  def __init__(self,bindaddr,statusfile='/tmp/status'):
      global MYCOLLECTOR
      MYCOLLECTOR=self
      # startup initialization
      self.bindaddr=bindaddr
      self.statusfile=statusfile
      # create a server
      self.s=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
      self.s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
      self.s.bind(self.bindaddr)
      self.s.listen(10)
  def main(self):
      global COLLECTOR_SERVER
      COLLECTOR_SERVER=self.bindaddr
      self.lastsave=time.mktime(time.localtime())
      self.pipe_r,self.pipe_w=os.pipe() # pipe for updates
      self.pid=os.fork()
      if self.pid==0: # main process
        # chroot
        from aglib import dochroot,sighup
        dochroot()
        signal.signal(signal.SIGHUP,sighup)
        signal.signal(signal.SIGUSR1,self.sigusr1)
        signal.signal(signal.SIGTERM,self.sigterm)
        if globals.daemon:
          globals.pidf.close()
        os.close(self.pipe_w)
        debug.dup()
        self.data['Start-at']=self.lastsave
        self.loadstats()
        self.data['Restart-at']=self.lastsave
        self.data['Hostname']=socket.gethostname()
        debug.echo(1,"collector(): service started, waiting for connections...")
        self.mypid=os.getpid()
        while 1:
          try:
            self.conn,addr=self.s.accept()
            socket_settimeout(self.conn,8)
            try:
              f=self.conn.makefile("rw",0)
              req=f.readline()
              debug.echo(5,"collector(): QUERY: ",req.strip())
              proto=0.0
              reg1=re.search("^GET /([^ ]*) HTTP/(1.[01])",req)
              if reg1:
                dirs=reg1.group(1).split("/")
                proto=float(reg1.group(2))
              else:
                reg1=re.search("^GET /([^ ]*)$",req.rstrip())
                if reg1:
                  proto=1.0
                  dirs=reg1.group(1).split("/")
              if proto>=1.0:
                args=self.parseargs(dirs[1:])
                lines={}
                while proto>=1.1:
                  line=f.readline()
                  if (line=="\r\n") or (line=="\n"):
                    break
                  a=line.split(": ")
                  if len(a)>1:
                    lines[a[0]]=a[1].strip()
                  else:
                    lines[a[0]]=a[0].strip()
                self.conn.sendall("HTTP/1.1 200 OK\r\n")
                self.conn.sendall("Content-Type: text/plain; charset=iso-8859-1\r\n")
                self.conn.sendall("\r\n");
                self.data["Uptime"]=uptime(self.data['Start-at'])
                self.data["Restart-age"]=uptime(self.data['Restart-at'])
                sstr=''
                if dirs[0]=="":
                  for key,value in self.data.items():
                    sstr+=key+":\t"+str(value)+"\r\n"
                  self.conn.sendall(sstr)
                elif dirs[0]=="mrtg":
                  for key in dirs[1:3]+['Uptime','Hostname']:
                    if key in self.data.keys():
                      try:
                        o="%-20.0f" % self.data[key]
                        sstr+=o.strip()+"\r\n"
                      except TypeError:
                        sstr+=str(self.data[key])+"\r\n"
                    else: # return 0 for unknown key
                      sstr+="0\r\n"
                  self.conn.sendall(sstr)
                else:
                  self.message("404 Not Found")
              else:
                self.message("501 Method Not Implemented")
                self.conn.sendall(req)
            except socket.error,(ec,es):
              if not ec in [32,104,107]:
                debug.echo(3,"collector(): socket.error: ",ec,es)
                debug.traceback(4,"collector(): ")
            except:
              debug.traceback(3,"collector(): ")
            try:
              self.conn.shutdown(2)
              self.conn.close()
            except socket.error:
              pass 
          except:
            try:
              self.conn.shutdown(2)
              self.conn.close()
            except:
              pass
        # while 1 ... end
      else:
        self.s.close()
        os.close(self.pipe_r)
      return self.pid
  def update(self,dirs):
      d=string.join(dirs,"/")+"\n"+' '*1024
      signal.signal(signal.SIGALRM,self.sigalrm)
      try:
        signal.alarm(2)
        os.write(self.pipe_w,d[:1024])
        signal.alarm(0)
      except OSError,(ec,es):
        signal.alarm(0)
        if ec==4: # Interrupted system call
          debug.echo(3,"collector(): writing to pipe: ",es)
      except:
        debug.echo(1,"collector(): unknown error on update")
        debug.traceback(3,"collector()")
      signal.signal(signal.SIGALRM,signal.SIG_DFL)
      os.kill(self.pid,signal.SIGUSR1)
  def sigalrm(self,sn,stack):
      # ignore this signal, only do an "interrupted system call"
      debug.echo(1,"collector(): WARNING: Interrupting write to pipe! (pipe full?)")
  def sigusr1(self,sn,stack):
      signal.signal(signal.SIGALRM,self.sigalrm)
      try:
        signal.alarm(2)
        sstr=os.read(self.pipe_r,1024)
        signal.alarm(0)
      except OSError,(ec,es):
        signal.alarm(0)
        if ec==4: # Interrupted system call
          debug.echo(3,"collector(): ",es)
          return
      signal.signal(signal.SIGALRM,signal.SIG_DFL)
      args=self.parseargs(sstr.split("/"))
      self.data["Childrens"]=0.0
      for k,v in args.items():
        try:
          self.data[k]+=float(v)
        except KeyError:
          self.data[k]=float(v)
        except ValueError:
          self.data[k]=v
        except TypeError:
          debug.echo(0,"collector(): TypeError: ",k,v,type(v))
      # save stats evry 5 minutes
      ctime=time.mktime(time.localtime())
      if (ctime-self.lastsave)>300:
        self.lastsave=ctime
        self.savestats()
  def sigterm(self,sn,stack):
      try:
        os.close(self.pipe_r)
        debug.echo(4,"collector(): Saving stats...")
        self.savestats()
      except:
        pass
      debug.echo(1,"collector(): Exiting - SIGTERM ...")
      os._exit(0)
  def loadstats(self):
      try:
        f=open(self.statusfile,'r')
        for line in f.readlines():
          a=line.split(":\t")
          try:
            self.data[a[0]]=float(a[1])
          except ValueError:
            self.data[a[0]]=a[1]
          except IndexError:
            # ignore wrong lines from status file
            pass
        f.close()
      except IOError,(ec,es):
        debug.echo(1,"collector(): loadstat error: ",es)
      except:
        debug.echo(1,"collector(): Can't load status file.")
        debug.traceback(1,"collector()")
  def savestats(self):
      f=open(self.statusfile,'w')
      for k,v in self.data.items():
        f.write(k+":\t"+str(v)+"\n")
      f.close()
      debug.echo(4,"collector(): Status saved.")
  def message(self,sstr):
      self.conn.sendall("HTTP/1.1 "+sstr+"\r\n\r\n")
  def parseargs(self,dirs):
      args={}
      for arg in dirs:
        a=arg.split("=")
        k=urllib.unquote_plus(a[0])
        if re.search("^[A-Za-z0-9-]*$",k):
          if len(a)>1:
            args[k]=urllib.unquote_plus(a[1])
          else:
            args[k]='1'
        else:
          debug.echo(1,"collector(): Unknown key: ",k)
      return args

def query(qt,args):
    if COLLECTOR_SERVER:
      s=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
      socket_settimeout(s,10)
      s.connect(COLLECTOR_SERVER)
      urlargv=''
      for arg in args:
        urlargv+="/"+urllib.quote_plus(arg,"/=")
      s.sendall("GET /"+qt+urlargv+" HTTP/1.1\r\n\r\n")
      f=s.makefile("r",0)
      lines=[]
      for _ in range(512):
        line=f.readline()
        if (line=="\r\n") or (line=="\n") or (line==""):
          break;
        lines.append(line)
      lines=[]
      s.shutdown(2)
      return s.recv(9999).rstrip()
    else:
      return None

def update(bs=-1,t=0,infected=0,tempfail=0):
    '''update statistics by 'bs' bytes, 't' time,
       if infected, then update also filtered bytes,files
    '''
    args=['Total-count','Total-bytes='+str(bs),'Total-time='+str(t),
          'Temp-fail='+str(tempfail),'Childrens=0']
    if infected:
      args.extend(['Filtered-count','Filtered-bytes='+str(bs)])
    try:
      MYCOLLECTOR.update(args+status.collect)
      status.collect=[]
    except:
      pass

class rrd:
  RRDDATA={
    'Total-bytes':		'totalbs',
    'Filtered-bytes':		'filteredbs',
    'Virus-bytes':		'virusbs',
    'Spam-bytes':		'spambs',
    'Total-count':		'total',
    'Filtered-count':		'filtered',
    'Virus-count':		'virus',
    'Spam-count':		'spam',
    'Temp-fail':		'fails',
    'Total-time':		'time'
  }
  DATES=[
    [3600*24,		'Day',		1,	600],
    [3600*24*7,		'Week',		6,	700],
    [3600*24*31,	'Month',	24,	775],
    [3600*24*365,	'Year',		288,	797]
  ]
  PARAMS='''--imgformat PNG --interlaced --width 576'''
  def __init__(self,path='./'):
      self.RRDPATH=path
      self.RRDFILE=os.path.join(path,'sagator.rrd')
  def create(self,xparams=''):
      # create an RRD
      cmd='rrdtool create %s %s' % (self.RRDFILE,xparams)
      for ds in self.RRDDATA.values():
        cmd+=' DS:%s:COUNTER:600:0:U' % ds
      for p1 in ['AVERAGE','MAX']:
        for p2 in self.DATES:
          cmd+=' RRA:%s:0.5:%d:%d' % (p1,p2[2],p2[3])
      os.system(cmd)
      return cmd
  def update(self):
      try:
        open(self.RRDFILE)
      except IOError:
        self.create()
      q=query('/',[])
      cmd='rrdtool update %s N' % self.RRDFILE
      for key,value in self.RRDDATA.items():
        reg1=re.compile('^%s:[ \t]*([0-9]+)' % key,re.M).search(q)
        if reg1:
          cmd+=':'+reg1.group(1)
        else:
          cmd+=':0'
          #print "WARNING: Key %s not found!" % key
      os.system(cmd)
      return self
  def system(self,s):
      os.system(s.replace('$STEP',self.STEP).replace('$PATH',self.RRDPATH))
  def range(self,delta,r,xparams):
      self.system('''
        rrdtool graph $PATH/%s-count.png %s %s \
        --vertical-label msgs/min --units-exponent 0 --start -%s \
        DEF:total=$PATH/sagator.rrd:total:AVERAGE \
        DEF:filtered=$PATH/sagator.rrd:filtered:AVERAGE \
        DEF:virus=$PATH/sagator.rrd:virus:AVERAGE \
        DEF:spam=$PATH/sagator.rrd:spam:AVERAGE \
        CDEF:mtotal=total,60,* \
        CDEF:mfiltered=filtered,60,* \
        CDEF:mvirus=virus,60,* \
        CDEF:mspam=spam,60,* \
        CDEF:dtotal=total,UN,0,total,IF,$STEP,* \
        CDEF:stotal=PREV,UN,dtotal,PREV,IF,dtotal,+ \
        CDEF:dfiltered=filtered,UN,0,filtered,IF,$STEP,* \
        CDEF:sfiltered=PREV,UN,dfiltered,PREV,IF,dfiltered,+ \
        CDEF:dvirus=virus,UN,0,virus,IF,$STEP,* \
        CDEF:svirus=PREV,UN,dvirus,PREV,IF,dvirus,+ \
        CDEF:dspam=spam,UN,0,spam,IF,$STEP,* \
        CDEF:sspam=PREV,UN,dspam,PREV,IF,dspam,+ \
        CDEF:rfiltered=sfiltered,100,*,stotal,/ \
        CDEF:rvirus=svirus,100,*,stotal,/ \
        CDEF:rspam=sspam,100,*,stotal,/ \
        AREA:mtotal#00a000:"Total   " \
        GPRINT:stotal:MAX:"%%8.0lf msgs" \
        GPRINT:mtotal:AVERAGE:"avg\\: %%5.2lf%%s/min" \
        GPRINT:mtotal:MAX:"max\\: %%5.2lf%%s/min\\l" \
        LINE2:mfiltered#c0c000:"Filtered" \
        GPRINT:sfiltered:MAX:"%%8.0lf msgs" \
        GPRINT:rfiltered:AVERAGE:"(%%6.2lf %%%%)\\l" \
        LINE2:mvirus#a00000:"Virus   " \
        GPRINT:svirus:MAX:"%%8.0lf msgs" \
        GPRINT:rvirus:AVERAGE:"(%%6.2lf %%%%)\\l" \
        LINE2:mspam#0000a0:"Spam    " \
        GPRINT:sspam:MAX:"%%8.0lf msgs" \
        GPRINT:rspam:AVERAGE:"(%%6.2lf %%%%)" \
        COMMENT:"                    Generated at %s by SAGATOR"
      ''' % (delta,self.PARAMS,xparams,r,time.strftime("%c")))
      self.system('''
        rrdtool graph $PATH/%s-bytes.png %s %s \
        --vertical-label bytes/min --start -%s \
        DEF:total=$PATH/sagator.rrd:totalbs:AVERAGE \
        DEF:filtered=$PATH/sagator.rrd:filteredbs:AVERAGE \
        DEF:virus=$PATH/sagator.rrd:virusbs:AVERAGE \
        DEF:spam=$PATH/sagator.rrd:spambs:AVERAGE \
        CDEF:mtotal=total,60,* \
        CDEF:mfiltered=filtered,60,* \
        CDEF:mvirus=virus,60,* \
        CDEF:mspam=spam,60,* \
        CDEF:dtotal=total,UN,0,total,IF,$STEP,* \
        CDEF:stotal=PREV,UN,dtotal,PREV,IF,dtotal,+ \
        CDEF:dfiltered=filtered,UN,0,filtered,IF,$STEP,* \
        CDEF:sfiltered=PREV,UN,dfiltered,PREV,IF,dfiltered,+ \
        CDEF:dvirus=virus,UN,0,virus,IF,$STEP,* \
        CDEF:svirus=PREV,UN,dvirus,PREV,IF,dvirus,+ \
        CDEF:dspam=spam,UN,0,spam,IF,$STEP,* \
        CDEF:sspam=PREV,UN,dspam,PREV,IF,dspam,+ \
        CDEF:rfiltered=sfiltered,100,*,stotal,/ \
        CDEF:rvirus=svirus,100,*,stotal,/ \
        CDEF:rspam=sspam,100,*,stotal,/ \
        AREA:mtotal#00a000:"Total   " \
        GPRINT:stotal:MAX:"%%8.2lf %%sB" \
        GPRINT:mtotal:AVERAGE:"avg\\: %%5.2lf %%sB/min" \
        GPRINT:mtotal:MAX:"max\\: %%5.2lf %%sB/min\\l" \
        LINE2:mfiltered#c0c000:"Filtered" \
        GPRINT:sfiltered:MAX:"%%8.2lf %%sB" \
        GPRINT:rfiltered:AVERAGE:"(%%6.2lf %%%%)\\l" \
        LINE2:mvirus#a00000:"Virus   " \
        GPRINT:svirus:MAX:"%%8.2lf %%sB" \
        GPRINT:rvirus:AVERAGE:"(%%6.2lf %%%%)\\l" \
        LINE2:mspam#0000a0:"Spam    " \
        GPRINT:sspam:MAX:"%%8.2lf %%sB" \
        GPRINT:rspam:AVERAGE:"(%%6.2lf %%%%)" \
        COMMENT:"                      Generated at %s by SAGATOR"
      ''' % (delta,self.PARAMS,xparams,r,time.strftime("%c")))
      self.system('''
        rrdtool graph $PATH/%s-time.png %s %s \
        --vertical-label "Sagator Load" --units-exponent 0 --start -%s \
        DEF:time=$PATH/sagator.rrd:time:MAX \
        AREA:time#00a000:"Sagator load" \
        GPRINT:time:AVERAGE:"avg\\: %%5.2lf" \
        GPRINT:time:MAX:"max\\: %%5.2lf"
      ''' % (delta,self.PARAMS,xparams,r))
  def graph(self,lazy='--lazy'):
      xparams=''
      for d in self.DATES:
        self.STEP=str(300*d[2])
        self.range(d[1].lower(),d[0],xparams)
        xparams=lazy
      return self

class status(interscan.match.match_any):
  '''
  status() interscanner
  
  Usage: status("String",scanner1(),scanner2(),...)
  
  Where: "String" is a string, which defines a prefix for status update
           This scring can't contain spaces.
         scanner1(),... are sagator's scanners
  
  This interscanner can be used to collect some other statistics.
  Preffered usage:
    SCANNERS=[
      status("Virus",virus_scanner1(),...),
      status("Spam",spam_scanner1(),...)
    ]
  This collects Virus-count,Virus-bytes,Spam-count,Spam-bytes in collector.
  '''
  name='' # making this scanner invisible
  collect=[]
  def __init__(self,string,*scanners):
      self.status_string=string
      interscan.match.match_any.__init__(self,scanners)
  def scanstream(self,stream):
      level,vir,ret=interscan.match.match_any.scanstream(self,stream)
      if vir:
        status.collect.extend([self.status_string+"-count",
          self.status_string+"-bytes="+str(len(stream))])
      return level,vir,ret

if __name__ == '__main__':
  if len(sys.argv)>1:
    # get host and port from collector service
    from config import SRV
    for service in SRV:
      if service.name=='collector()':
        COLLECTOR_SERVER=service.BINDTO
    if sys.argv[1]=='rrdtool':
      rrd(sys.argv[2]).update().graph()
    elif sys.argv[1]=='rrdupdate':
      rrd(sys.argv[2]).update()
    elif sys.argv[1]=='rrdgraph':
      rrd(sys.argv[2]).graph('')
    else:
      print query(sys.argv[1],sys.argv[2:])
