JeevesLib.py

"""API for Python Jeeves libary.

  :synopsis: Functions for creating sensitive values, labels, and policies.

.. moduleauthor:: Travis Hance <tjhance7@gmail.com>
.. moduleauthor:: Jean Yang <jeanyang@csail.mit.edu>

"""

from env.VarEnv import VarEnv
from env.PolicyEnv import PolicyEnv
from env.PathVars import PathVars
from env.WritePolicyEnv import WritePolicyEnv
from smt.Z3 import Z3
from fast.AST import Facet, fexpr_cast, Constant, Var, Not, FExpr, Unassigned, FObject, jeevesState
from eval.Eval import partialEval
import copy

def init():
  """Initialization function for Jeeves library.

  You should always call this before you do anything Jeeves-y.

  """
  jeevesState.init()

  # TODO this needs to be GC'ed somehow

def supports_jeeves(f):
  f.__jeeves = 0
  return f

@supports_jeeves
def mkLabel(varName = "", uniquify=True):
  """Makes a label to associate with policies and sensitive values.

  :param varName: Optional variable name (to help with debugging).
  :type varName: string
  :returns: Var - fresh label.
  """
  label = jeevesState.policyenv.mkLabel(varName, uniquify)
  jeevesState.all_labels[label.name] = label
  return label

@supports_jeeves
def doesLabelExist(varName):
  return varName in jeevesState.all_labels

@supports_jeeves
def getLabel(varName):
  return jeevesState.all_labels[varName]

@supports_jeeves
def restrict(varLabel, pred, use_empty_env=False):
  """Associates a policy with a label.

  :param varLabel: Label to associate with policy.
  :type varLabel: string
  :param pred: Policy: function taking output channel and returning Boolean result.
  :type pred: T -> bool, where T is the type of the output channel
  """
  jeevesState.policyenv.restrict(varLabel, pred, use_empty_env)

@supports_jeeves
def mkSensitive(varLabel, vHigh, vLow):
  """Creates a sensitive value with two facets.

  :param varLabel: Label to associate with sensitive value.
  :type varLabel: Var
  :param vHigh: High-confidentiality facet for viewers with restricted access.
  :type vHigh: T
  :param vLow: Low-confidentiality facet for other viewers.
  :type vLow: T
  """

  if isinstance(varLabel, Var):
    return Facet(varLabel, fexpr_cast(vHigh), fexpr_cast(vLow))
  else:
    return JeevesLib.jif(varLabel, lambda:vHigh, lambda:vLow)

@supports_jeeves
def concretize(ctxt, v):
  """Projects out a single value to the viewer.

  :param ctxt: Output channel (viewer).
  :type ctxt: T, where policies have type T -> bool
  :param v: Value to concretize.
  :type v: FExpr
  :returns: The concrete (non-faceted) version of T under the policies in the environment.
  """
  return jeevesState.policyenv.concretizeExp(ctxt, v, jeevesState.pathenv.getEnv())

@supports_jeeves
def jif(cond, thn_fn, els_fn):
  condTrans = partialEval(fexpr_cast(cond), jeevesState.pathenv.getEnv())
  if condTrans.type != bool:
    raise TypeError("jif must take a boolean as a condition")
  return jif2(condTrans, thn_fn, els_fn)

def jif2(cond, thn_fn, els_fn):
  if isinstance(cond, Constant):
    return thn_fn() if cond.v else els_fn()

  elif isinstance(cond, Facet):
    if not isinstance(cond.cond, Var):
      raise TypeError("facet conditional is of type %s"
                      % cond.cond.__class__.__name__)

    with PositiveVariable(cond.cond):
      thn = jif2(cond.thn, thn_fn, els_fn)
    with NegativeVariable(cond.cond):
      els = jif2(cond.els, thn_fn, els_fn)

    return Facet(cond.cond, thn, els)

  else:
    raise TypeError("jif condition must be a constant or a var")

# supports short-circuiting
# without short-circuiting jif is unnecessary
# are there performance issues?
@supports_jeeves
def jand(l, r): # inputs are functions
  left = l()
  if not isinstance(left, FExpr):
    return left and r()
  return jif(left, r, lambda:left)

@supports_jeeves
def jor(l, r): # inputs are functions
  left = l()
  if not isinstance(left, FExpr):
    return left or r()
  return jif(left, lambda:left, r)

# this one is more straightforward
# just takes an expression
@supports_jeeves
def jnot(f):
  if isinstance(f, FExpr):
    return Not(f)
  else:
    return not f

@supports_jeeves
def jassign(old, new, base_env={}):
  res = new
  for vs in jeevesState.pathenv.conditions:
    (var, val) = (vs.var, vs.val)
    if var.name not in base_env:
      if val:
        res = Facet(var, res, old)
      else:
        res = Facet(var, old, res)
  if isinstance(res, FExpr):
    return partialEval(res, {}, True)
  else:
    return res

'''
@supports_jeeves
def jhasElt(lst, f):
  acc = False
  # Short circuits.
  for elt in lst:
    isElt = f(elt) # TODO: This should eventually be japply of f to elt.
    if isinstance(isElt, FExpr):
      acc = jor(lambda: isElt, lambda: acc)
    else:
      if isElt:
        return True
  return acc 

@supports_jeeves
def jhas(lst, v):
  return jhasElt(lst, lambda x: x == v)
'''

class PositiveVariable:
  def __init__(self, var):
    self.var = var
  def __enter__(self):
    jeevesState.pathenv.push(self.var, True)
  def __exit__(self, type, value, traceback):
    jeevesState.pathenv.pop()

class NegativeVariable:
  def __init__(self, var):
    self.var = var
  def __enter__(self):
    jeevesState.pathenv.push(self.var, False)
  def __exit__(self, type, value, traceback):
    jeevesState.pathenv.pop()

def liftTuple(t):
  t = fexpr_cast(t)
  if isinstance(t, FObject):
    return t.v
  elif isinstance(t, Facet):
    a = liftTuple(t.thn)
    b = liftTuple(t.els)
    return tuple([Facet(t.cond, a1, b1) for (a1, b1) in zip(a, b)])
  else:
    raise TypeError("bad use of liftTuple")

class Namespace:
  def __init__(self, kw, funcname):
    self.__dict__.update(kw)
    self.__dict__['_jeeves_funcname'] = funcname
    self.__dict__['_jeeves_base_env'] = jeevesState.pathenv.getEnv()

  def __setattr__(self, attr, value):
    self.__dict__[attr] = jassign(self.__dict__.get(attr, Unassigned("variable '%s' in %s" % (attr, self._jeeves_funcname))), value, self.__dict__['_jeeves_base_env'])

@supports_jeeves
def jgetattr(obj, attr):
  if isinstance(obj, FExpr):
    return getattr(obj, attr)
  else:
    return getattr(obj, attr) if hasattr(obj, attr) else Unassigned("attribute '%s'" % attr)

@supports_jeeves
def jgetitem(obj, item):
  try:
    return obj[item]
  except (KeyError, KeyError, TypeError) as e:
    return Unassigned("item '%s'" % attr)

@supports_jeeves
def jmap(iterable, mapper):
  if isinstance(iterable, JList2):
    return jmap_jlist2(iterable, mapper)
  if isinstance(iterable, FObject) and isinstance(iterable.v, JList2):
    return jmap_jlist2(iterable.v, mapper)

  iterable = partialEval(fexpr_cast(iterable), jeevesState.pathenv.getEnv())
  return FObject(JList(jmap2(iterable, mapper)))
def jmap2(iterator, mapper):
  if isinstance(iterator, Facet):
    if jeevesState.pathenv.hasPosVar(iterator.cond):
      return jmap2(iterator.thn, mapper)
    if jeevesState.pathenv.hasNegVar(iterator.cond):
      return jmap2(iterator.els, mapper)
    with PositiveVariable(iterator.cond):
      thn = jmap2(iterator.thn, mapper)
    with NegativeVariable(iterator.cond):
      els = jmap2(iterator.els, mapper)
    return Facet(iterator.cond, thn, els)
  elif isinstance(iterator, FObject):
    return jmap2(iterator.v, mapper)
  elif isinstance(iterator, JList):
    return jmap2(iterator.l, mapper)
  elif isinstance(iterator, JList2):
    return jmap2(iterator.convert_to_jlist1().l, mapper)
  elif isinstance(iterator, list) or isinstance(iterator, tuple):
    return FObject([mapper(item) for item in iterator])
  else:
    return jmap2(iterator.__iter__(), mapper)

def jmap_jlist2(jlist2, mapper):
  ans = JList2([])
  env = jeevesState.pathenv.getEnv()
  for i, e in jlist2.l:
    popcount = 0
    for vname, vval in e.iteritems():
      if vname not in env:
        v = getLabel(vname)
        jeevesState.pathenv.push(v, vval)
        popcount += 1
      elif env[vname] != vval:
        break
    else:
      ans.l.append((mapper(i), e))
    for _ in xrange(popcount):
      jeevesState.pathenv.pop()
  return FObject(ans)

def facetMapper(facet, fn, wrapper=fexpr_cast):
  if isinstance(facet, Facet):
    return Facet(facet.cond, facetMapper(facet.thn, fn, wrapper), facetMapper(facet.els, fn, wrapper))
  elif isinstance(facet, Constant) or isinstance(facet, FObject):
    return wrapper(fn(facet.v))

class JList:
  def validate(self):
    def foo(x):
      assert isinstance(x, list), 'thingy is ' + str(x.l.v)
      return x
    facetMapper(self.l, foo, lambda x : x)

  def __init__(self, l):
    self.l = l if isinstance(l, FExpr) else FObject(l)
    self.validate()
  def __getitem__(self, i):
    return self.l[i]
  def __setitem__(self, i, val):
    self.l[i] = jassign(self.l[i], val)

  def __len__(self):
    return self.l.__len__()
  def __iter__(self):
    return self.l.__iter__()

  def append(self, val):
    l2 = facetMapper(self.l, list, FObject) #deep copy
    l2.append(val)
    self.l = jassign(self.l, l2)
    self.validate()

  def prettyPrint(self):
    def tryPrint(x):
      return x.__class__.__name__
      '''
      try:
        return x.__class__.__name__ #x.prettyPrint()
      except AttributeError:
        return str(x)
      '''
    return str(len(self.l)) #''.join(map(tryPrint, self.l))

class JList2:
  def __init__(self, l=[]):
    if isinstance(l, list):
      self.l = [(i, {}) for i in l]
    else:
      raise NotImplementedError
  
  def append(self, val):
    self.l.append((val, jeevesState.pathenv.getEnv()))

  def eval(self, env):
    return [i for i,e in self.l if all(env[getLabel(v)] == e[v] for v in e)]

  def vars(self):
    all_vars = set()
    for _, e in self.l:
      all_vars.update(set(e.keys()))
    return {getLabel(v) for v in all_vars}

  def convert_to_jlist1(self):
    all_vars = [v.name for v in self.vars()]
    def rec(cur_e, i):
      if i == len(all_vars):
        return FObject([i for i,e in self.l if all(cur_e[v] == e[v] for v in e)])
      else:
        cur_e1 = dict(cur_e)
        cur_e2 = dict(cur_e)
        cur_e1[all_vars[i]] = True
        cur_e2[all_vars[i]] = False
        return Facet(getLabel(all_vars[i]),
            rec(cur_e1, i+1), rec(cur_e2, i+1))
    return JList(rec({}, 0))

  def __getitem__(self, i):
    return self.convert_to_jlist1().__getitem__(i)

  def __setitem__(self, i, val):
    raise NotImplementedError

  def __len__(self):
    return self.convert_to_jlist1().__len__()

class JIterator:
  def __init__(self, l):
    self.l = l

@supports_jeeves
def jfun(f, *args, **kw):
  if hasattr(f, '__jeeves'):
    return f(*args, **kw)
  else:
    env = jeevesState.pathenv.getEnv()
    if len(args) > 0:
      return jfun2(f, args, kw, 0, partialEval(fexpr_cast(args[0]), env), [])
    else:
      it = kw.__iter__()
      try:
        fst = next(it)
      except StopIteration:
        return fexpr_cast(f())
      return jfun3(f, kw, it, fst, partialEval(fexpr_cast(kw[fst]), env), (), {})

def jfun2(f, args, kw, i, arg, args_concrete):
  if isinstance(arg, Constant) or isinstance(arg, FObject):
    env = jeevesState.pathenv.getEnv()
    if i < len(args) - 1:
      return jfun2(f, args, kw, i+1, partialEval(fexpr_cast(args[i+1]), env), tuple(list(args_concrete) + [arg.v]))
    else:
      it = kw.__iter__()
      try:
        fst = next(it)
      except StopIteration:
        return fexpr_cast(f(*tuple(list(args_concrete) + [arg.v])))
      return jfun3(f, kw, it, fst, partialEval(fexpr_cast(kw[fst]), env), tuple(list(args_concrete) + [arg.v]), {})
  else:
    with PositiveVariable(arg.cond):
      thn = jfun2(f, args, kw, i, arg.thn, args_concrete)
    with NegativeVariable(arg.cond):
      els = jfun2(f, args, kw, i, arg.els, args_concrete)
    return Facet(arg.cond, thn, els)

from itertools import tee
def jfun3(f, kw, it, key, val, args_concrete, kw_concrete):
  if isinstance(val, Constant) or isinstance(val, FObject):
    kw_c = dict(kw_concrete)
    kw_c[key] = val.v
    try:
      next_key = next(it)
    except StopIteration:
      return fexpr_cast(f(*args_concrete, **kw_c))
    env = jeevesState.pathenv.getEnv()
    return jfun3(f, kw, it, next_key, partialEval(fexpr_cast(kw[next_key]), env), args_concrete, kw_c)
  else:
    it1, it2 = tee(it)
    with PositiveVariable(val.cond):
      thn = jfun3(f, kw, it1, key, val.thn, args_concrete, kw_concrete)
    with NegativeVariable(val.cond):
      els = jfun3(f, kw, it2, key, val.els, args_concrete, kw_concrete)
    return Facet(val.cond, thn, els)

def evalToConcrete(f):
    g = partialEval(fexpr_cast(f), jeevesState.pathenv.getEnv())
    if isinstance(g, Constant):
      return g.v
    elif isinstance(g, FObject):
      return g.v
    else:
      raise Exception("wow such error: evalToConcrete on non-concrete thingy-ma-bob")

from jlib.JContainer import *