'''
Sagators statistics collector and distributor, version 0.3.1

(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):
      global MYCOLLECTOR
      MYCOLLECTOR=self
      # startup initialization
      self.bindaddr=bindaddr
      # 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(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)
        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)
        str=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(str.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('/tmp/status','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]
        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('/tmp/status','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,str):
      self.conn.sendall("HTTP/1.1 "+str+"\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=[]
      str=''
      str+=s.recv(9999).rstrip()
      s.shutdown(2)
      return str
    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 status(interscan.match.match_any):
  '''
  status() interscanner
  
  Usage: status("String",scanner1(),scanner2(),...)
  
  Where: "String" is a string, which defines a prefix for status update
         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.HOST,service.PORT)
    print query(sys.argv[1],sys.argv[2:])
