#!/usr/bin/env python

"""
Run the queries described in README

cmd: env ABAC_CN=1 keystore=`pwd` ./query.py

"""

import os
import sys
import ABAC
import time
import datetime
from datetime import timedelta

debug=0

ctxt = ABAC.Context()

cred_count = 2 + 2 * #VAL#

def get_msec(e_time) :
    msec_delta=0
    if( int(e_time.seconds) !=0 ) :
        msec_delta= int(e_time.seconds) *1000
    if( int(e_time.microseconds) !=0) :
        msec_delta = msec_delta + int(e_time.microseconds)/1000
    return msec_delta

def get_micro(e_time) :
    micro_delta=0
    if( int(e_time.seconds) !=0 ) :
        micro_delta= int(e_time.seconds) *1000000
    if( int(e_time.microseconds) !=0) :
        micro_delta = micro_delta + int(e_time.microseconds)
    return micro_delta

def extract_delta(starttime, endtime) :
    """ given a start time, and a endtime, extract delta """
    elapsed_time = (endtime - starttime)
# Only handle in seconds/microseconds
    if ( int(elapsed_time.days) != 0 ) :
        sys.stderr,write("%s is longer than a day !!!" % msg)
        exit(1)
    return elapsed_time


# Keystore is the directory containing the principal credentials.
# Load existing principals and/or policy credentials
if (os.environ.has_key("keystore")) :
    keystore=os.environ["keystore"]
    starttime = datetime.datetime.now()
    ctxt.load_directory(keystore)
    endtime = datetime.datetime.now()
    elapsed_load=extract_delta(starttime, endtime)
    elapsed_msec=get_msec(elapsed_load)
    sys.stderr.write("%d %d LOAD(msec)\n" % (cred_count,elapsed_msec))
else:
    print("keystore is not set...")
    exit(1)

##########################################################################
# dump the loaded principals/policies
#

fd=os.open("creds_dump",os.O_WRONLY|os.O_CREAT)
out = ctxt.context_principals()
for x in out[1]:
    os.write(fd, x.string())
    os.write(fd,"\n")
out = ctxt.context_credentials()
for c in out[1]:
    string="%s <- %s" % (c.head_string(), c.tail_string())
    os.write(fd,string) 
    os.write(fd,"\n")
os.close(fd)

##########################################################################
# Does JohnX likes John0 ?
# role = [keyid:JohnX].role:after 
# p [Keyid:john0]
def goodQuery() :
    aid="John%s_ID.pem"% #VAL#
    aID=ABAC.ID(aid)
    a=aID.id_keyid()

    bID=ABAC.ID("John0_ID.pem")
    b=bID.id_keyid()

    role = ABAC.Role(a,"likes")
    p = ABAC.Role(b)

    print "\n===good============ johnN.likes <- john0 "
    starttime = datetime.datetime.now()
    out = ctxt.query(role, p)
    endtime = datetime.datetime.now()
    if(debug):
        print "good query start-> %s\n" % starttime
        print "good query end  -> %s\n" % endtime
    for c in out[1]:
        print "%s <- %s" % (c.head_string(), c.tail_string())
    return extract_delta(starttime, endtime)
    

##########################################################################
# Does John0 likes JohnX ?
# role = [keyid:JohnX].role:after 
# p [Keyid:john0]
def badQuery() :
    bid="John%s_ID.pem"% #VAL#
    bID=ABAC.ID(bid)
    b=bID.id_keyid()

    aID=ABAC.ID("John0_ID.pem")
    a=aID.id_keyid()
    
    role = ABAC.Role(a,"likes")
    p = ABAC.Role(b)
    
    print "\n===bad============ john0.likes <- johnX "
    starttime = datetime.datetime.now()
    out = ctxt.query(role, p)
    endtime = datetime.datetime.now()
    for c in out[1]:
        print "%s <- %s" % (c.head_string(), c.tail_string())
    return extract_delta(starttime, endtime)
        
##############################################################

#skip the first one
e_time=goodQuery()
elapsed_micro=get_micro(e_time)
sys.stderr.write("%d %d GOOD_1(micro)\n" % (cred_count,elapsed_micro))

elapsed_good=timedelta()
k=10
while(k):
    e_time=goodQuery()
    elapsed_good=elapsed_good+e_time
    k=k-1
    if(k==9):
       elapsed_micro=get_micro(e_time)
       sys.stderr.write("%d %d GOOD_2(micro)\n" % (cred_count,elapsed_micro))

    if(debug):
        elapsed_micro=get_micro(e_time)
        sys.stderr.write("%d %d GOOD_%d(micro)\n" % (cred_count,elapsed_micro,k))
elapsed_micro=get_micro(elapsed_good)
sys.stderr.write("%d %d GOOD_t(micro)\n" % (cred_count,elapsed_micro/10))


###############################################################

e_time=badQuery()
elapsed_micro=get_micro(e_time)
sys.stderr.write("%d %d BAD_1(micro)\n" % (cred_count,elapsed_micro))

elapsed_bad=timedelta()
k=10
while(k):
    e_time=badQuery()
    elapsed_bad=elapsed_bad+e_time
    k=k-1
    if(k==9):
        elapsed_micro=get_micro(e_time)
        sys.stderr.write("%d %d BAD_2(micro)\n" % (cred_count,elapsed_micro))

    if(debug):
        elapsed_micro=get_micro(e_time)
        sys.stderr.write("%d %d BAD_%d(micro)\n" % (cred_count,elapsed_micro,k))

elapsed_micro=get_micro(elapsed_bad)
sys.stderr.write("%d %d BAD_t(micro)\n" % (cred_count,elapsed_micro/10))