#! /usr/bin/env python

# Copyright 1999-2009 University of Chicago
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
# http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Usage Stats Collector

Receives packets from a UDP socket and stores them in the local file system.
The files are written in a subdirectory of the packet directory, named by the
date this program received the packet. The packet files themselves are named
by the hour which the collector received the packet.

Note that the date and hour are relative to UTC.

The following command-line options are available:
  -h                       Print this message and exit
  -p PORT                  Listen for usage packets on UDP port PORT
                           [default: 4810]

"""

import sys
import SocketServer
import os
import getopt
import socket
import signal
import thread
import time

from globus.usage.packetfile import PacketFile
import globus.usage.usageconfig as usageconfig

collector_port = "4810"

config = usageconfig.UsageConfig()

packetfile = PacketFile(config.data_directory, "a")
packetfile_lock = thread.allocate_lock()

pidfile_path = None
receiver = None

def shutdownthread():
    if receiver != None:
        receiver.shutdown()
    sys.exit(0)
    
def termhandler(signum, frame):
    if pidfile_path != None:
        os.remove(pidfile_path)
    thread.start_new_thread(shutdownthread, ())

def packetcloser():
    last_close = int(time.time() / 60 / 60) * 60 * 60
    while True:
        checktime = int(time.time() / 60 / 60) * 60 * 60
        if checktime >= (last_close + 60 * 60):
            last_close = checktime
            packetfile_lock.acquire()
            packetfile.close()
            packetfile_lock.release()
        nextcheck = checktime + 60 * 60
        time.sleep(nextcheck - time.time())
        

signal.signal(signal.SIGTERM, termhandler)
closer = thread.start_new_thread(packetcloser, ())

class Receiver(SocketServer.BaseRequestHandler):
    """
    Usage stats receiver class

    Receives packets on its UDP port and writes them to a file
    """

    def __init__(self, request, client_address, base_server):
        SocketServer.BaseRequestHandler.__init__(self, request, client_address, base_server)

    def handle(self):
        """
        Handle UDP packet
        """
        addr = socket.inet_aton(self.client_address[0])
        packetfile_lock.acquire()
        packetfile.write_packet(addr, self.request[0])
        packetfile_lock.release()

def usage(exe):
    """
    Return usage string
    """
    usagestr = "Usage: " + os.path.basename(exe) + " [options]"
    usagestr += """
Options:
  -h                       Print this message and exit
  -b                       Run collector as a background process
  -d DIRECTORY             Write packets in DIRECTORY instead of the default
  -f FILE                  Write process id of backgrounded process to FILE
  -p PORT                  Listen for usage packets on UDP port PORT"""
    print usagestr 

def main(argv):
    """
    globus-usage-collector main program

    """
    global pidfile_path
    global receiver

    background = 0
    pid = -1

    try:
        port = int(collector_port)
        opts, args = getopt.getopt(argv[1:], "d:p:hbf:")
        for option_name in opts:
            if option_name[0] == '-h':
                usage(argv[0])
                sys.exit(0)
            elif option_name[0] == "-p":
                port = int(option_name[1])
            elif option_name[0] == "-d":
                config.data_directory = option_name[1]
            elif option_name[0] == "-b":
                background = 1
            elif option_name[0] == "-f":
                pidfile_path = option_name[1]

        if args != []:
            print "Unexpected command-line: ", args
            return 1

        if background != 0:
            pid = os.fork()
            if pid < 0:
                print "Error forking process"
                sys.exit(1)
            elif pid > 0:
                print "Running in background (%d)" % (pid)
                if pidfile_path != None:
                    pidfile = file(pidfile_path, "w")
                    pidfile.write("%d\n" % (pid))
                    pidfile.close()
                sys.exit(0)

        receiver = SocketServer.UDPServer(('', port), Receiver)
        print "Starting usage collector on UDP port " + str(port)
        sys.stdin.close()
        sys.stdout.close()
        sys.stderr.close()
        sys.stdin = file("/dev/null", "r")
        sys.stdout = file("/dev/null", "a")
        sys.stderr = file("/dev/null", "a")
        receiver.serve_forever()
    except KeyboardInterrupt:
        print "Quitting..."
        return 0
    except getopt.GetoptError, detail:
        print "Error: " + str(detail)
        return 1

if __name__ == "__main__":
    __exit_code__ = main(sys.argv)
    sys.exit(__exit_code__)
# vim: set ts=4:sw=4:syntax=python
