taint

commit 5953f0609a707a31c9866134aaea99cbee656300
parent 59aea60d87769e5c6a5495b1feb55dfb6ac62e7c
Author: nolash <dev@holbrook.no>
Date:   Fri, 16 Apr 2021 16:55:02 +0200

Invoke result handler on account tag action

Diffstat:
M
taint/cache.py
 | 
11
+++++++++--
M
taint/tag.py
 | 
1
-
M
taint/taint.py
 | 
8
++++++++
M
tests/eth/test_syncer.py
 | 
21
++++++++++++++++++---

4 files changed, 35 insertions(+), 6 deletions(-)
diff --git a/taint/cache.py b/taint/cache.py
@@ -184,10 +184,15 @@ class Cache(Salter):
         return (subjects, objects)
 
 
+    def add_account(self, account, label):
+        self.cache_bloom.add_raw(account.account, label)
+
+
     def add_subject(self, account):
         if not isinstance(account, Account):
             raise TypeError('subject must be type crypto_account_cache.account.Account')
-        self.cache_bloom.add_raw(account.account, 'subject')
+        #self.cache_bloom.add_raw(account.account, 'subject')
+        self.add_account(account, 'subject')
         logg.debug('added subject {}'.format(account))
         self.subjects[account.account] = account
 
@@ -195,7 +200,8 @@ class Cache(Salter):
     def add_object(self, account):
         if not isinstance(account, Account):
             raise TypeError('subject must be type crypto_account_cache.account.Account')
-        self.cache_bloom.add_raw(account.account, 'object')
+        #self.cache_bloom.add_raw(account.account, 'object')
+        self.add_account(account, 'object')
         logg.debug('added object {}'.format(account))
         self.objects[account.account] = account
 
@@ -217,6 +223,7 @@ class Cache(Salter):
 
         # TODO: watch out, this currently scales geometrically
         (subjects, objects) = self.divide(accounts)
+        logg.debug('subjects {} objects {}'.format(subjects, objects))
         for subject in subjects:
             for objct in objects:
                 subject.connect(objct) 
diff --git a/taint/tag.py b/taint/tag.py
@@ -20,7 +20,6 @@ class TagPool:
         self.pool[z].append(account)
         
         logg.debug('pool register account {} to tag sum {}'.format(account, z.hex()))
-        
 
         return True
 
diff --git a/taint/taint.py b/taint/taint.py
@@ -19,6 +19,12 @@ class Tainter(Cache):
             self.store.save(self.cache_bloom.serialize())
 
 
+    def add_account(self, account, label):
+        super(Tainter, self).add_account(account, label)
+        if self.result_handler != None:
+            self.result_handler.register(account)
+
+
     def filter(self, conn, block, tx, storer):
         for output in tx.outputs:
             for inpt in tx.inputs:
@@ -34,3 +40,5 @@ class Tainter(Cache):
                 if self.result_handler != None:
                     for account in subjects:
                         self.result_handler.register(account)
+                    for account in objects:
+                        self.result_handler.register(account)
diff --git a/tests/eth/test_syncer.py b/tests/eth/test_syncer.py
@@ -15,7 +15,10 @@ from hexathon import strip_0x
 # local imports
 from taint.taint import Tainter
 from taint.account import Account
-from taint.tag import TagPool
+from taint.tag import (
+        TagPool,
+        Tag,
+        )
 
 
 class TestSyncer(EthTesterCase):
@@ -28,15 +31,17 @@ class TestSyncer(EthTesterCase):
         self.alice_account = bytes.fromhex(strip_0x(self.accounts[0]))
         self.alice = Account(self.chain_spec, self.alice_account, label='alice')
         self.alice.tag(b'foo')
-        self.alice.tag(b'bar')
 
+        self.bob_account = bytes.fromhex(strip_0x(self.accounts[1]))
+        self.bob = Account(self.chain_spec, self.bob_account, label='bob')
+        self.bob.tag(b'bar')
 
         self.filter_bits_size = 1024
         self.result_handler = TagPool()
         self.frontend = Tainter(self.chain_spec, self.filter_bits_size, result_handler=self.result_handler)
+        self.frontend.add_subject(self.bob)
         self.frontend.add_subject(self.alice)
 
-
         self.syncer.add_filter(self.frontend)
 
 
@@ -55,6 +60,16 @@ class TestSyncer(EthTesterCase):
 
         self.assertTrue(self.frontend.have(4, 0))
 
+        tag = Tag()
+        tag.create(b'foo')
+        tag.create(b'bar')
+        z = tag.get()
+
+        self.assertEqual(len(self.result_handler.pool.keys()), 3)
+        self.assertIn(z, self.result_handler.pool.keys())
+        self.assertIn(self.alice, self.result_handler.pool[z])
+        self.assertIn(self.bob, self.result_handler.pool[z])
+
 
 if __name__ == '__main__':
     unittest.main()