fixes and improvements in slither-simil

utils/similarity/__main__.py

@@ -34,21 +34,16 @@ def parse_args():
                         dest='filename',
                         help='contract.sol')
 
-    parser.add_argument('--contract',
+    parser.add_argument('--fname',
                         action='store',
-                        dest='contract',
-                        help='Contract')
+                        dest='fname',
+                        help='Target function')
 
-    parser.add_argument('--filter',
+    parser.add_argument('--ext',
                         action='store',
-                        dest='filter',
+                        dest='ext',
                         help='Extension to filter contracts')
 
-    parser.add_argument('--fname',
-                        action='store',
-                        dest='fname',
-                        help='Function name')
-
     parser.add_argument('--nsamples',
                         action='store',
                         type=int,

utils/similarity/encode.py

@@ -13,10 +13,24 @@ from slither.solc_parsing.variables.local_variable import *
 from slither.solc_parsing.variables.local_variable_init_from_tuple import *
 
 from .cache import load_cache
+from crytic_compile.platform.solc import InvalidCompilation
+
+simil_logger = logging.getLogger("Slither-simil")
+compiler_logger = logging.getLogger("CryticCompile")
+compiler_logger.setLevel(logging.CRITICAL)
+slither_logger = logging.getLogger("Slither")
+slither_logger.setLevel(logging.CRITICAL)
+
+def parse_target(target):
+    parts = target.split('.')
+    if len(parts) == 1:
+        return None, parts[0]
+    elif len(parts) == 2:
+        return parts
+    else:
+        simil_logger.error("Invalid target. It should be 'function' or 'Contract.function'")
 
-logger = logging.getLogger("Slither-simil")
-
-def load_and_encode(infile, model, ext=None, nsamples=None, **kwargs):
+def load_and_encode(infile, model, filter=None, nsamples=None, **kwargs):
     r = dict()
     if infile.endswith(".npz"):
         r = load_cache(infile, nsamples=nsamples)
@@ -30,7 +44,7 @@ def load_and_encode(infile, model, ext=None, nsamples=None, **kwargs):
 
     return r
 
-def load_contracts(dirname, ext=None, nsamples=None):
+def load_contracts(dirname, ext=None, nsamples=None, **kwargs):
     r = []
     walk = list(os.walk(dirname))
     for x, y, files in walk:
@@ -160,17 +174,17 @@ def encode_ir(ir):
 
     # default
     else:
-        logger.error(type(ir),"is missing encoding!")
+        simil_logger.error(type(ir),"is missing encoding!")
         return ''
  
-def encode_contract(filename, **kwargs):
+def encode_contract(cfilename, **kwargs):
     r = dict()
 
     # Init slither
     try:
-        slither = Slither(filename, **kwargs)
+        slither = Slither(cfilename, **kwargs)
     except:
-        logger.error("Compilation failed")
+        simil_logger.error("Compilation failed for %s using %s", cfilename, kwargs['solc'])
         return r
 
     # Iterate over all the contracts
@@ -185,7 +199,7 @@ def encode_contract(filename, **kwargs):
                 if function.nodes == []:
                     continue
 
-                x = (filename,contract.name,function.name) 
+                x = (cfilename,contract.name,function.name) 
 
                 r[x] = []
 

utils/similarity/info.py

@@ -4,7 +4,7 @@ import os.path
 import traceback
 
 from fastText import load_model
-from .encode import encode_contract
+from .encode import parse_target, encode_contract
 
 logging.basicConfig()
 logger = logging.getLogger("Slither-simil")
@@ -20,20 +20,20 @@ def info(args):
             model = None
 
         filename = args.filename
-        contract = args.contract
+        contract, fname = parse_target(args.fname)
         solc = args.solc
-        fname = args.fname
+        
         if filename is None and contract is None and fname is None:
-            print(args.model,"uses the following words:")
+            logger.info("%s uses the following words:",args.model)
             for word in model.get_words():
-                print(word)
+                logger.info(word)
             sys.exit(0)
 
         if filename is None or contract is None or fname is None:
             logger.error('The encode mode requires filename, contract and fname parameters.')
             sys.exit(-1)
 
-        irs = encode_contract(filename, solc=solc)
+        irs = encode_contract(filename, **vars(args))
         if len(irs) == 0:
             sys.exit(-1)
         

utils/similarity/plot.py

@@ -13,7 +13,7 @@ except ImportError:
     plt = None
     
 from fastText import load_model
-from .encode import load_and_encode
+from .encode import load_and_encode, parse_target
 
 logger = logging.getLogger("Slither-simil")
 
@@ -29,19 +29,19 @@ def plot(args):
         model = args.model
         model = load_model(model)
         filename = args.filename
-        contract = args.contract
-        fname = args.fname
-        solc = args.solc
+        #contract = args.contract
+        contract, fname = parse_target(args.fname)
+        #solc = args.solc
         infile = args.input
-        ext = args.filter
-        nsamples = args.nsamples
+        #ext = args.filter
+        #nsamples = args.nsamples
 
         if fname is None or infile is None:
             logger.error('The plot mode requieres fname and input parameters.')
             sys.exit(-1)
 
         logger.info('Loading data..')
-        cache = load_and_encode(infile, model, ext=ext, solc=solc, nsamples=nsamples)
+        cache = load_and_encode(infile, **vars(args))
 
         data = list()
         fs = list()

utils/similarity/test.py

@@ -6,8 +6,8 @@ import operator
 import numpy as np
 
 from fastText import load_model
-from .encode import encode_contract, load_and_encode
-from .cache import load_cache, save_cache
+from .encode import encode_contract, load_and_encode, parse_target
+from .cache import save_cache
 from .similarity import similarity
 
 logger = logging.getLogger("Slither-simil")
@@ -18,25 +18,22 @@ def test(args):
         model = args.model
         model = load_model(model)
         filename = args.filename
-        contract = args.contract
-        fname = args.fname
-        solc = args.solc
+        contract, fname = parse_target(args.fname) 
         infile = args.input
-        ext = args.filter
         ntop = args.ntop
 
         if filename is None or contract is None or fname is None or infile is None:
             logger.error('The test mode requires filename, contract, fname and input parameters.')
             sys.exit(-1)
 
-        irs = encode_contract(filename,solc=solc)
+        irs = encode_contract(filename, **vars(args))
         if len(irs) == 0:
             sys.exit(-1)
 
         y = " ".join(irs[(filename,contract,fname)])
         
         fvector = model.get_sentence_vector(y)
-        cache = load_and_encode(infile, model, ext=ext, solc=solc)
+        cache = load_and_encode(infile, **vars(args))
         #save_cache("cache.npz", cache)
 
         r = dict()

utils/similarity/train.py

@@ -16,22 +16,20 @@ def train(args):
     try:
         last_data_train_filename = "last_data_train.txt"
         model_filename = args.model
-        solc = args.solc
         dirname = args.input
-        ext = args.filter
         nsamples = args.nsamples
 
         if dirname is None:
             logger.error('The train mode requires the input parameter.')
             sys.exit(-1)
 
-        contracts = load_contracts(dirname, ext=ext, nsamples=nsamples)
+        contracts = load_contracts(dirname, **vars(args))
         logger.info('Saving extracted data into %s', last_data_train_filename)
         cache = []
         with open(last_data_train_filename, 'w') as f:
             for filename in contracts:
                 #cache[filename] = dict()
-                for (filename, contract, function), ir in encode_contract(filename,solc).items():
+                for (filename, contract, function), ir in encode_contract(filename, **vars(args)).items():
                     if ir != []:
                         x = " ".join(ir)
                         f.write(x+"\n")