'''
    Implementing Erasure Policies Using Taint Analysis
    Copyright (C) 2010 Filippo Del Tedesco and Alejandro Russo

    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 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
'''
from datetime import datetime

__all__ = ['erasure_source','erasure_escape', 'erasure', 'lazy_erasure', 'retain']

###################
##ERASURE WRAPPER##
###################
class Erasure:
   '''The erasure wrapper'''
   def __init__(self, value):
      self.val = value
      self.tstamps=set()

   def erase(self):
      if isinstance(self.val,str):
         self.val = ""
      elif isinstance(self.val,int):
         self.val= 0
      self.tstamps=set()

   def __getattr__(self, name):
      if name in ['__doc__',           # Documentation
                  '__repr__',          # Called to show strings on the screen
                  '__hash__',          # Used to create dictionaries
                  '__class__',         # To indicate that it is an string
                  '__setattr__',       # Because this should behave as an string (no attribute)
                  '__delattr__',       # Because this should behave as an string (no attribute)
                  '__str__']:
        return getattr(self.val, name)
      elif name == '__radd__':
         def _shadow(other):
            obj=Erasure(other)
            return obj.__add__(self)
         return _shadow
      else: 
         return method_proxy(getattr(self.val,name), True, self.tstamps)
##END##############


###################
###LIBRARY'S STATE#
###################
###PART1: Erasure redefinition of standard classes
STR = Erasure
INT = Erasure
eclasses = {str : STR, int : INT }
###PART2: Erasure state
class Table:
   dependencies = {}
   policies     = []
   
   def compute_closure_ts(self, o):
      current=extract_ts(o)
      return current 

   def erase_with_ts(self, ts):
      for t in ts: 
         if t in self.dependencies and any([f(t) for f in self.policies]):
            for o in self.dependencies[t]:
               o.erase()
            del self.dependencies[t]

   def add_dependency(self,o,t):
      if t in self.dependencies:
         self.dependencies[t].append(o)
      else:
         self.dependencies[t]=[o]

   def add_dependencies(self,o):
      def inner(o):
         for t in o.tstamps:
            self.add_dependency(o,t)
      mapt(o,inner,lambda o: hasattr(o,'tstamps'))

   def lazy_erasure(self, o = None):
      if o is None:
         raise KeyError
      elif callable(o):
         self.policies.append(o)
      else:
         ts = self.compute_closure_ts(o)
         self.policies.append(lambda t: t in ts)

   def eager_erasure(self, o=None):
      temp=self.policies
      if o is None:
         ts=self.dependencies.key()
      else:
         ts=self.compute_closure_ts(o)
      self.policies=[lambda t: t in ts]
      self.erase_with_ts(ts)
      self.policies=temp

estate=Table()


#######################
###LIBRARY'S INTERFACE#
#######################
def erasure_source(f):
   ''' Makes the return object of f an erasure-aware object'''
   def inner(*args,**kwargs):
      localS=set([])
      ff = method_proxy(f,True,localS)
      r  = ff(*args, **kwargs)
      ts = datetime.today()
      r.tstamps.add(ts)
      estate.add_dependency(r,ts)
      return r
   return inner

def erasure_escape(f):
   '''
   Checks that all its arguments are erasure-approved before calling f.
   Otherwise an erasure action is triggered
   '''
   def inner(*args,**kwargs):
      ts=set([])
      for a in args:
         ts.update(estate.compute_closure_ts(a))
      for k,v in kwargs:
         ts.update(estate.compute_closure_ts(v))
      estate.erase_with_ts(ts)
      return f(*args,**kwargs)
   return inner

def retain(f):
   '''Unmount the erasure structure for the return object of f'''
   def inner(*args,**kwargs):
      ff=method_proxy(f,False,set([]))
      return ff(*args,**kwargs)
   return inner

def erasure(obj=None):
   '''Implement eager erasure'''
   estate.eager_erasure(obj)

def lazy_erasure(p1,p2=None):
   '''Implement lazy erasure 
      USAGE -lerasure(policy)
            -lerasure(object)
            -lerasure(object,policy)
   '''
   if not(p2==None):  
      def policy(time):
         if time in p1.tstamps and p2(time):
            return True
         else:
            return False
      estate.lazy_erasure(policy)
   else:
      estate.lazy_erasure(p1)   


#######################
##AUXILIARY FUNCTIONS##
#######################
def mapt(o, f, check=lambda o: type(o) in eclasses.keys()):
   '''
   Applies "f" on true cases for "check"
   by structural induction on o
   '''
   if not isinstance(o,Erasure) and not(o):
      return o
   if check(o):
      return f(o)
   elif isinstance(o, list):
      return [mapt(x, f, check) for x in o]
   elif isinstance(o, tuple):
      return tuple([mapt(x, f, check) for x in o])
   elif isinstance(o, set):
      return  set([mapt(x, f, check) for x in o])
   elif isinstance(o, dict):
      klass = type(o) 
      return klass([(k, mapt(v, f, check)) for k, v in o.items()])
   else:
      return o

def extract_ts(o):
   '''Copies the tstamps of o into ts'''
   ts=set([])
   mapt(o,lambda o: ts.update(o.tstamps),lambda o: hasattr(o,'tstamps'))
   return ts

def inject_ts(o,ts):
   '''Copies the taints of ts into obj.tstamps'''
   mapt(o, lambda o: o.tstamps.update(ts), lambda o: hasattr(o, 'tstamps'))

def unwrap(o):
   '''Turns an erasure-aware object in a conventional one'''
   return mapt(o, lambda o: o.val , lambda o:hasattr(o,'val'))

def wrap(o, ts):
   '''Turns a conventional object in an erasure aware one'''
   def erclass(o):
      '''Erasurizer'''
      klass = type(o)
      if klass in eclasses.keys():
         return eclasses[klass](o)
      else:
         raise KeyError
   o=mapt(o, erclass)
   inject_ts(o, ts)
   return o

def method_proxy(m,w,ts):
   '''
   It proxies the capability of a method by unwrapping arguments and wrapping results
   m: the method we are proxying
   w: should we wrap the result?
   ts: the time stamps we need to consider in the final wrapping
   '''
   def inner(*args,**kargs):
      temp_ts = ts.copy()  
      temp_ts.update(extract_ts(args))
      temp_ts.update(extract_ts(kargs))
      r=m(*unwrap(args),**unwrap(kargs))
      if w:
         r=wrap(r,temp_ts)
         estate.add_dependencies(r)
      return r
   return inner