#!/usr/bin/env python
"""
Tests the local FroNtier squid
"""
#
# Assumes:
#
#   1) environmental variables SAME_OK and SAME_ERROR are defined
#   2) environmental variable $CMS_PATH is defined
#   3) file $CMS_PATH/SITECONF/local/JobConfig/site-local-config.xml
#      contains the location of the local FroNtier squid server
#
__revision__ = "$Id: test_squid.py,v 1.7 2007/03/01 16:07:53 ericw Exp $"
import os
import sys
import urllib2
from xml.dom.minidom import parseString
from xml.dom.minidom import parse
import base64
import zlib 
import curses.ascii
import time
#
# Print out node name
#
print "node " + os.uname()[1] 
#
# Check that environmental variable SAME_OK is set
#
if not os.environ.has_key("SAME_OK"):
    print "test_squid.py: Error. SAME_OK not defined"
    sys.exit(1)
same_ok = int(os.environ["SAME_OK"])    
#
# Check that environmental variable SAME_ERROR is set
#
if not os.environ.has_key("SAME_ERROR"):
    print "test_squid.py: Error. SAME_ERROR not defined"
    sys.exit(1)
same_error = int(os.environ["SAME_ERROR"])    
#
# Check that environmental variable CMS_PATH is set
#
if not os.environ.has_key("CMS_PATH"):
    print "test_squid.py: Error. CMS_PATH not defined"
    sys.exit(same_error)
#
# Check that file $CMS_PATH/SITECONF/local/JobConfig/site-local-config.xml
# exists
#
slcfil = os.environ["CMS_PATH"] + \
         "/SITECONF/local/JobConfig/site-local-config.xml"
if not os.path.exists(slcfil):
    print "test_squid.py: Error. file " + slcfil + " does not exist"
    sys.exit(same_error)
#
# Read and parse site-local-config.xml into a dom
# See http://docs.python.org/lib/module-xml.dom.minidom.html
#
slcdom = parse(slcfil)
#
# Work out site name from site-local-config.xml
#
silist = slcdom.getElementsByTagName("site")
if len(silist) == 0:
    site = "UNKNOWN"
else:
    stag = silist[0]
    site = stag.getAttribute("name")
print "site " + site    
#
# Work out local FroNtier squid server from site-local-config.xml
# Note that http_proxy is not set when running on CERN site
#
if site != "CERN":
#
# Check for at least one proxy tag
#
    prlist = slcdom.getElementsByTagName("proxy")
    if len(prlist) == 0:
        print "test_squid.py: Error. no proxy tag in file " + slcfil
        sys.exit(same_error)    
#
# Work out address of squid
#
    prtag = prlist[0]
    squid = prtag.getAttribute("url")
#
# Set proxy server
#
    print "squid " + squid   
    os.environ["http_proxy"] = squid
#
# Following code from Sinisa Veseli 
#
# Set parameters
#
frontierUrl = "http://cmsfrontier.cern.ch:8080/FrontierInt/Frontier"
frontierQuery = "SELECT 1 FROM DUAL"
decodeFlag = True
refreshFlag = True
statsFlag = False
retrieveZiplevel = ""
#
# Print parameters 
#  
if statsFlag:
    pass
else:
    print "Using Frontier URL: ", frontierUrl
    print "Query: ", frontierQuery
    print "Decode results: ", decodeFlag
    print "Refresh cache: ", refreshFlag
#
# Encode query
#
compQuery = zlib.compress(frontierQuery, 9)
encQuery = base64.binascii.b2a_base64(compQuery).replace("+", ".")
#
# Set up FroNtier request
#
format = "%s?type=frontier_request:1:DEFAULT&encoding=BLOB%s&p1=%s"
frontierRequest = format % (frontierUrl, retrieveZiplevel, encQuery)
if statsFlag:
    pass
else:
    print "\nFrontier Request:\n", frontierRequest 
#
# Add refresh header if needed
#
request = urllib2.Request(frontierRequest)
if refreshFlag:
    request.add_header("pragma", "no-cache")
#
# Start and time query
#
queryStart = time.localtime()
if statsFlag:
    pass
else:
    print "\nQuery started: ", time.strftime("%m/%d/%y %H:%M:%S %Z", queryStart)
t1 = time.time()
#
try:
    result = urllib2.urlopen(request).read()
except urllib2.URLError, e:
    rc, message = e.reason
    if site != "CERN":
        print "test_squid.py: Error. squid " + squid + " is down, "
        print "                      unreachable or will not reply."
        print "                      " + message
        sys.exit(same_error)
    else:
        print "test_squid.py: Error. FroNtier server " + frontierUrl + \
              " is down, "
        print "                      unreachable or will not reply."
        print "                      " + message
        sys.exit(same_error)
except:
    print "test_squid.py: Error."
    sys.exit(same_error)
#
t2 = time.time()
queryEnd = time.localtime()
if statsFlag:
    duration = (t2-t1)
    size = len(result)
    print duration, size, size/duration
else:
    print "Query ended: ", time.strftime("%m/%d/%y %H:%M:%S %Z", queryEnd)
    print "Query time: %0.2f [seconds]\n" % (t2-t1)
#
# Print out result
#
if decodeFlag:
    print "Query result:\n", result
    dom = parseString(result)
    dataList = dom.getElementsByTagName("data")
    # Control characters represent records, but I won't bother with that now,
    # and will simply replace those by space.
    for data in dataList:
        if data.firstChild is not None:
            row = base64.decodestring(data.firstChild.data)
            if retrieveZiplevel != "":
                row = zlib.decompress(row)
            control = [ '\x00', '\x01', '\x02', '\x03', '\x04',
                        '\x05', '\x06', '\x08', '\x09', '\x0a',
                        '\x0b', '\x0c', '\x0d', '\x1b', '\x17'  ]    
            for c in control:
                row = row.replace(c, ' ')
            print "\nFields: "
            endFirstRow = row.find('\x07')
            firstRow = row[:endFirstRow]
            for c in firstRow:
                if curses.ascii.isctrl(c):
                    firstRow = firstRow.replace(c, '\n')
            print firstRow
            print "\nRecords:"
            pos = endFirstRow + 1
            while True:
                newrow = row[pos:]
                endRow = newrow.find('\x07')
                if endRow < 0:
                    break
                fixedRow = newrow[:endRow]
                pos = pos + endRow + 1
                fixedRow = fixedRow.replace('\n', '')
                print fixedRow
#
# Exit with success exit code
#
print "OK"
sys.exit(same_ok)