diff --git a/smalisca/modules/module_sql_models.py b/smalisca/modules/module_sql_models.py index 157d645..60f94d4 100644 --- a/smalisca/modules/module_sql_models.py +++ b/smalisca/modules/module_sql_models.py @@ -32,11 +32,16 @@ """Represent an App as SQL data""" +import re import textwrap +import sqlite3 import sqlalchemy as sql -from sqlalchemy import ForeignKey +from sqlalchemy import ForeignKey, event from sqlalchemy.orm import relationship, scoped_session, sessionmaker from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy import event +from sqlalchemy.engine import Engine +import time import smalisca.core.smalisca_config as config from smalisca.core.smalisca_logging import log @@ -45,28 +50,70 @@ Base = declarative_base() +import logging +logging.basicConfig() +log = logging.getLogger(__file__) +log.setLevel(logging.DEBUG) + # Tables defining relationships between entities # Classes <-> Properties class_properties_table = sql.Table( 'class_properties', Base.metadata, sql.Column('class_id', sql.Integer, ForeignKey('classes.id')), - sql.Column('prop_id', sql.Integer, ForeignKey('properties.id')) + sql.Column('prop_id', sql.Integer, ForeignKey('properties.id')), + mysql_charset='utf8' ) # Classes <-> Const Strings class_const_strings_table = sql.Table( 'class_const_strings', Base.metadata, sql.Column('class_id', sql.Integer, ForeignKey('classes.id')), - sql.Column('const_string_id', sql.Integer, ForeignKey('const_strings.id')) + sql.Column('const_string_id', sql.Integer, ForeignKey('const_strings.id')), + mysql_charset='utf8' ) # Classes <-> Methods class_methods_table = sql.Table( 'class_methods', Base.metadata, sql.Column('class_id', sql.Integer, ForeignKey('classes.id')), - sql.Column('method_id', sql.Integer, ForeignKey('methods.id')) + sql.Column('method_id', sql.Integer, ForeignKey('methods.id')), + mysql_charset='utf8' ) +""" +# U N C O M M E N T T H I S I F Y O U W A N T T O M O N I T O R +# T I M E E X E C U T I O N O F A L L T H E Q U E R I E S +# WARNING: This is very VERBOSE!! +#============================================================================== +@event.listens_for(Engine, "before_cursor_execute") +def before_cursor_execute(conn, cursor, statement, + parameters, context, executemany): + conn.info.setdefault('query_start_time', []).append(time.time()) + log.debug("Start Query: %s", statement) + +@event.listens_for(Engine, "after_cursor_execute") +def after_cursor_execute(conn, cursor, statement, + parameters, context, executemany): + total = time.time() - conn.info['query_start_time'].pop(-1) + log.debug("Query Complete!") + log.debug("Total Time: %f", total) + +#============================================================================== +""" + +def regexp(pattern, input): + try: + """ + re.I = ignore case + re.X verbose mode, this enables to write comments inside regex expr. + e.g: r'(.*) # some comment that will be ignored' + """ + reg = re.compile(pattern) + except re.error: + #log.debug("Failed to compile the following regex's pattern >>> {0}".format(pattern)) + return False + return reg.search(input) is not None + class SmaliClass(Base): """Models a Smali class @@ -84,19 +131,13 @@ class SmaliClass(Base): """ __tablename__ = "classes" - # Constraints - __table_args__ = ( - sql.UniqueConstraint( - 'class_name', 'class_type', 'depth', 'path', name='unique_class'), - ) - # Fields id = sql.Column(sql.Integer, primary_key=True) - class_name = sql.Column(sql.Text) - class_type = sql.Column(sql.Text) - class_package = sql.Column(sql.Text) + class_name = sql.Column(sql.Text(collation='utf8_general_ci'), unique=False) + class_type = sql.Column(sql.Text(collation='utf8_general_ci'), unique=False) + class_package = sql.Column(sql.Text(collation='utf8_general_ci'), unique=False) depth = sql.Column(sql.Integer) - path = sql.Column(sql.Text) + path = sql.Column(sql.String(255), unique=False) # Relationships properties = relationship( @@ -119,6 +160,11 @@ def to_string(self): ) return textwrap.dedent(s) + def toDict(self): + return {'class_name' : self.class_name, + 'class_type' : self.class_type, + 'class_depth' : self.depth} + def __str__(self): return self.to_string() @@ -134,26 +180,23 @@ class SmaliProperty(Base): property_name (str): Name of the property property_type (str): Property type (e.g. Ljava/lang/String)# property_info (str): Additional property info (private, static, finale, etc.) - property_class (str): The class this property belongs to + class_name (str): The class this property belongs to """ __tablename__ = "properties" - # Constraints - __table_args__ = ( - sql.UniqueConstraint( - 'property_name', 'property_type', - 'property_info', 'property_class', - name='unique_property' - ), - ) - # Fields id = sql.Column(sql.Integer, primary_key=True) - property_name = sql.Column(sql.Text) - property_type = sql.Column(sql.Text) - property_info = sql.Column(sql.Text) - property_class = sql.Column(sql.Text) + property_name = sql.Column(sql.Text(collation='utf8_general_ci')) + property_type = sql.Column(sql.Text(collation='utf8_general_ci')) + property_info = sql.Column(sql.Text(collation='utf8_general_ci')) + class_name = sql.Column(sql.Text(collation='utf8_general_ci')) + + def toDict(self): + return {'field_name' : self.property_name, + 'field_type' : self.property_type, + 'field_info' : self.property_info, + 'class_name' : self.class_name} def to_string(self): s = """ @@ -163,7 +206,7 @@ def to_string(self): \t[+] Info: \t%s \t[+] Class: \t%s """ % (self.id, self.property_name, self.property_type, - self.property_info, self.property_class) + self.property_info, self.class_name) return textwrap.dedent(s) def __str__(self): @@ -186,9 +229,14 @@ class SmaliConstString(Base): # Fields id = sql.Column(sql.Integer, primary_key=True) - const_string_var = sql.Column(sql.Text) - const_string_value = sql.Column(sql.Text) - const_string_class = sql.Column(sql.Text) + const_string_var = sql.Column(sql.Text(collation='utf8_general_ci')) + const_string_value = sql.Column(sql.Text(collation='utf8_general_ci')) + class_name = sql.Column(sql.Text(collation='utf8_general_ci')) + + def toDict(self): + return {'const_string_var' : self.const_string_var, + 'const_string_value' : self.const_string_value, + 'class_name' : self.class_name} def to_string(self): s = """ @@ -196,7 +244,7 @@ def to_string(self): \t[+] Variable: \t%s \t[+] Value: \t%s \t[+] Class: \t%s - """ % (self.id, self.const_string_var, self.const_string_value, self.const_string_class) + """ % (self.id, self.const_string_var, self.const_string_value, self.class_name) return textwrap.dedent(s) @@ -216,26 +264,25 @@ class SmaliMethod(Base): method_type (str): Method type (public, abstract, constructor) method_args (str): Method arguments (e.g. Landroid/os/Parcelable;Ljava/lang/ClassLoader;) method_ret (str): Methods return value (Z, I, [I, etc.) - method_class (str): The class the method belongs to + class_name (str): The class the method belongs to """ __tablename__ = "methods" - # Constraints - __table_args__ = ( - sql.UniqueConstraint( - 'method_name', 'method_type', 'method_args', - 'method_ret', 'method_class', - name='unique_method' - ), - ) - # Fields id = sql.Column(sql.Integer, primary_key=True) - method_name = sql.Column(sql.Text) - method_type = sql.Column(sql.Text) - method_args = sql.Column(sql.Text) - method_ret = sql.Column(sql.Text) - method_class = sql.Column(sql.Text) + method_name = sql.Column(sql.Text(collation='utf8_general_ci')) + method_type = sql.Column(sql.String(collation='utf8_general_ci', length=255)) + method_args = sql.Column(sql.Text(collation='utf8_general_ci')) + method_ret = sql.Column(sql.Text(collation='utf8_general_ci')) + class_name = sql.Column(sql.Text(collation='utf8_general_ci')) + + def toDict(self): + return {'method_name' : self.method_name, + 'method_type' : self.method_type, + 'method_args' : self.method_args, + 'method_ret' : self.method_ret, + 'class_name' : self.method_ret} + def to_string(self): s = """ @@ -246,7 +293,7 @@ def to_string(self): \t[+] Ret: \t%s \t[+] Class: \t%s """ % (self.id, self.method_name, self.method_type, - self.method_args, self.method_ret, self.method_class) + self.method_args, self.method_ret, self.class_name) return textwrap.dedent(s) def __str__(self): @@ -276,17 +323,26 @@ class SmaliCall(Base): id = sql.Column(sql.Integer, primary_key=True) # Source class/method/args - from_class = sql.Column(sql.Text) - from_method = sql.Column(sql.Text) - local_args = sql.Column(sql.Text) + from_class = sql.Column(sql.Text(collation='utf8_general_ci')) + from_method = sql.Column(sql.Text(collation='utf8_general_ci')) + local_args = sql.Column(sql.Text(collation='utf8_general_ci')) # Destination class/method/args - dst_class = sql.Column(sql.Text) - dst_method = sql.Column(sql.Text) - dst_args = sql.Column(sql.Text) + dst_class = sql.Column(sql.Text(collation='utf8_general_ci')) + dst_method = sql.Column(sql.Text(collation='utf8_general_ci')) + dst_args = sql.Column(sql.Text(collation='utf8_general_ci')) # Return value - ret = sql.Column(sql.Text) + ret = sql.Column(sql.Text(collation='utf8_general_ci')) + + def toDict(self): + return {'from_class' : self.from_class, + 'from_method' : self.from_method, + 'local_args' : self.local_args, + 'dst_class' : self.dst_class, + 'dst_method' : self.dst_method, + 'dst_args' : self.dst_args, + 'ret' : self.ret} # FIXME: Add prettytable def to_string(self): @@ -331,9 +387,24 @@ class AppSQLModel: methods (list): List of methods calls (list): List of calls + How to use: + DATABASE='mysql://{0}:{1}@0.0.0.0:3306/{2}?charset=utf8mb4'.format(user, password, mysql_database_name) + myDB = AppSQLModel(DATABASE) + [...] + myDB.add_classes(classes) + myDB.add_properties(properties) + [...] + + """ - def __init__(self, sqlitedb): + @classmethod + def connection(cls, raw_con, connection_record): + if 'sqlite' in type(raw_con).__module__: + raw_con.create_function('regexp', 2, regexp) + + + def __init__(self, db): """Init the app SQL model Args: @@ -343,7 +414,10 @@ def __init__(self, sqlitedb): AppSqlModel: Instance of AppSQLModel """ - self.engine = sql.create_engine('sqlite:///' + sqlitedb) + #self.engine = sql.create_engine('sqlite:///' + sqlitedb) + self.engine = sql.create_engine(db, encoding="utf-8", convert_unicode=False) + event.listen(self.engine, 'connect', AppSQLModel.connection) + Base.metadata.create_all(self.engine) # Create session @@ -358,6 +432,7 @@ def __init__(self, sqlitedb): self.properties = [] self.methods = [] self.calls = [] + self.models_stack = [] def get_class_by_name(self, classname): """Returns class obj specified by name @@ -375,11 +450,21 @@ class object # Check if any results: try: if self.db.query(class_obj.exists()): - return class_obj.one() + return class_obj.first() except sql.orm.exc.NoResultFound: log.warn("No result found") return None + except sql.orm.exc.MultipleResultsFound: + log.error("MultipleResultsFound:\n\tClass = {0}".format(classname)) + for q in class_obj.all(): + log.debug("\t+ {0}".format(q.class_name)) + raise sql.orm.exc.MultipleResultsFound('MultipleResultsFound') + except: + log.error("get_class_by_name failed") + for q in class_obj.all(): + log.debug("\t+ {0}".format(q.class_name)) + raise def get_classes(self): """Returns all classes @@ -426,14 +511,23 @@ def get_calls(self): """ return self.db.query(SmaliCall).all() - def add_class(self, class_obj): + def save_all(self): + self.db.add_all(self.models_stack) + del self.models_stack[:] + + + def add_classes(self, classes): + self.engine.execute(SmaliClass.__table__.insert(), + classes) + + def add_class(self, class_obj, stack=False): """Add new class Args: class_obj (dict): Class object to insert """ - log.debug(class_obj) + #log.debug(class_obj) new_class = SmaliClass( class_name=class_obj['name'], class_type=class_obj['type'], @@ -444,9 +538,17 @@ def add_class(self, class_obj): # Add new class self.classes[class_obj['name']] = new_class - self.db.merge(new_class) + # self.db.merge(new_class, dont_load=True) + if stack is False: + self.db.add(new_class) + else: self.models_stack.append(new_class) + - def add_property(self, prop): + def add_properties(self, properties): + self.engine.execute(SmaliProperty.__table__.insert(), + properties) + + def add_property(self, prop, stack=False): """Adds property to class Args: @@ -458,7 +560,7 @@ def add_property(self, prop): property_name=prop['name'], property_type=prop['type'], property_info=prop['info'], - property_class=prop['class'] + class_name=prop['class'] ) # Append new property to class @@ -466,13 +568,21 @@ def add_property(self, prop): # Add to DB try: - self.db.merge(class_obj) - + # self.db.merge(class_obj, dont_load=True) + if stack is False: + self.db.add(class_obj) + else: self.models_stack.append(class_obj) except sql.exc.IntegrityError: self.db.rollback() - log.error("Found NOT unique values") + log.error("Insertion Failed: Found NOT unique values for property: {0}".format(new_prop)) + raise sql.exc.IntegrityError("Insertion Failed") + + + def add_const_strings(self, const_strings): + self.engine.execute(SmaliConstString.__table__.insert(), + const_strings) - def add_const_string(self, const_string): + def add_const_string(self, const_string, stack=False): """Adds const string to class Args: @@ -483,7 +593,7 @@ def add_const_string(self, const_string): new_const_string = SmaliConstString( const_string_var=const_string['name'], const_string_value=const_string['value'], - const_string_class=const_string['class'] + class_name=const_string['class'] ) # Append new const-string to class @@ -491,13 +601,20 @@ def add_const_string(self, const_string): # Add to DB try: - self.db.merge(class_obj) - + # self.db.merge(class_obj) + if stack is False: + self.db.add(class_obj) + else: self.models_stack.append(class_obj) except: self.db.rollback() - log.error("Failed inserting const-string\:%s" % new_const_string) + log.error("Insertion Failed: Failed inserting const-string\:%s" % new_const_string) + raise + + def add_methods(self, methods): + self.engine.execute(SmaliMethod.__table__.insert(), + methods) - def add_method(self, method): + def add_method(self, method, stack=False): """Adds property to class Args: @@ -510,7 +627,7 @@ def add_method(self, method): method_type=method['type'], method_args=method['args'], method_ret=method['return'], - method_class=method['class'] + class_name=method['class'] ) # Append new method to class @@ -518,13 +635,21 @@ def add_method(self, method): # Add to DB try: - self.db.merge(class_obj) + # self.db.merge(class_obj) + if stack is False: + self.db.add(class_obj) + else: self.models_stack.append(class_obj) except sql.exc.IntegrityError: self.db.rollback() - log.error("Found NOT unique values") + log.error("Insertion Failed: Found NOT unique values for method: {0}".format(new_method)) + raise sql.exc.IntegrityError("Insertion Failed") + + def add_calls(self, calls): + self.engine.execute(SmaliCall.__table__.insert(), + calls) - def add_call(self, call): + def add_call(self, call, stack=False): """Adds calls to class Args: @@ -549,7 +674,10 @@ def add_call(self, call): ) # Add new call to DB - self.db.merge(new_call) + # self.db.merge(new_call) + if stack is False: + self.db.add(new_call) + else: self.models_stack.append(new_call) def get_session(self): """Returns DB session