diff --git a/android_db.py b/android_db.py index 50ae2dc..e97e233 100644 --- a/android_db.py +++ b/android_db.py @@ -3,9 +3,8 @@ import re import csv import sys csv.field_size_limit(2**31-1) -param_str = re.compile(r'(?<=\()[ ,\?]+(?=\))') -validator = re.compile(r'^\?(?: *\, *\?)*$') -param_sep = re.compile(r' *\, *') +dict_params = re.compile(r':([^\ \,\.\\\(\)\=]+)') +list_params = re.compile(r'\?') class AndroidSQLConn: def __init__(self,device,filepath,use_root=True): self.device = device @@ -22,21 +21,24 @@ class AndroidSQLConn: else: return str(param) def _sub_params_(SQL,params): - p_string = param_str.search(SQL).group(0) - if not validator.match(p_string): - raise ValueError('Invalid substitution') + params = iter(params) + try: + return list_params.sub(lambda match: AndroidSQLConn._quote_param_(next(params)),SQL) + except StopIteration: + raise TypeError("Not enough parameters") + + def _sub_params_dict_(SQL,params): + try: + return dict_params.sub(lambda match: AndroidSQLConn._quote_param_(params[match.group(1)]),SQL) + except KeyError: + raise TypeError("Parameter specified but not in mapping") - n_params = len(param_sep.split(p_string)) - if len(params) < n_params: - raise ValueError('Not enough parameters supplied') - new_str = ','.join(map( - AndroidSQLConn._quote_param_,params[:n_params] - )) - print(*map(repr,(new_str,SQL))) - return param_str.sub(re.escape(new_str),SQL,1) def execute(self,SQL,params = None): if params: - SQL = AndroidSQLConn._sub_params_(SQL,params) + if isinstance(params,dict): + SQL = AndroidSQLConn._sub_params_dict_(SQL,params) + else: + SQL = AndroidSQLConn._sub_params_(SQL,params) SQL = shlex.quote(SQL) if self.root: shell = self.device.sudo @@ -45,3 +47,6 @@ class AndroidSQLConn: out = shell('sqlite3','-header','-csv',shlex.quote(self.filepath),SQL,output="out") if out: return csv.DictReader(out.splitlines()) + +if __name__ == "__main__": + print(AndroidSQLConn._sub_params_("SELECT * FROM table WHERE name=?,age=?;",["boby;DROP TABLE table",1])) \ No newline at end of file