rl-institut NESP2 database.py sql injection

@@ -3,7 +3,7 @@
 import pandas as pd
 from sqlalchemy import create_engine
 from sqlalchemy.orm import scoped_session, sessionmaker
-import geoalchemy2
+from sqlalchemy.sql import text
 from sqlalchemy import Table, MetaData
 from sqlalchemy.ext.declarative import declarative_base
 import random
@@ -75,53 +75,71 @@ def filter_materialized_view(
         limit=None,
         keys=None,
 ):
+    """
+
+    :param engine:
+    :param view_name: name of the view in the database (NOT A USER INPUT)
+    :param schema: name of the schema in the database (NOT A USER INPUT)
+    :param state_code: admin code of nigerian state (PROOFED USER INPUT)
+    :param area: boundaries for settlement's area filter (USER INPUT)
+    :param distance_grid: boundaries for settlement's distance to grid filter (USER INPUT)
+    :param building: boundaries for settlement's building count filter (USER INPUT)
+    :param buildingfp: boundaries for settlement's building percentage of area filter (USER INPUT)
+    :param limit: boundaries for settlement area filter (NOT A USER INPUT)
+    :param keys: list of columns to query values from (NOT A USER INPUT)
+    :return: returned data from the query
+    """
+
+    # to hold query parameters
+    values = {}
+
     if schema is not None:
         view_name = "{}.{}".format(schema, view_name)
+
     if limit is None:
         limit = ""
     else:
-        limit = " LIMIT {}".format(limit)
+        values["limit"] = int(limit)
+        limit = " LIMIT :limit"
 
-    filter_cond = ""
+    filter_cond = []
 
     if state_code is not None:
         key = "adm1_pcode"
-        filter_cond = f" WHERE {view_name}.{key}='{state_code}'"
+        filter_cond += [f"{view_name}.{key} = :{key}"]
+        values[key] = state_code
 
     if area is not None:
         key = "area_km2"
-        if "WHERE" in filter_cond:
-            filter_cond = filter_cond + f" AND {view_name}.{key} >= {area[0]} AND" \
-                                        f" {view_name}.{key} <= {area[1]}"
-        else:
-            filter_cond = f" WHERE {view_name}.{key} >= {area[0]} AND {view_name}.{key} <= {area[1]}"
+        val1 = key + "_1"
+        val2 = key + "_2"
+        filter_cond += [f"{view_name}.{key} >= :{val1}", f"{view_name}.{key} <= :{val2}"]
+        values[val1] = float(area[0])
+        values[val2] = float(area[1])
 
     if distance_grid is not None:
         key = "grid_dist_km"
-        if "WHERE" in filter_cond:
-            filter_cond = filter_cond + f" AND {view_name}.{key} >= {distance_grid[0]} AND" \
-                                        f" {view_name}.{key} <= {distance_grid[1]}"
-        else:
-            filter_cond = f" WHERE {view_name}.{key} >= {distance_grid[0]} AND" \
-                          f" {view_name}.{key} <= {distance_grid[1]}"
+        val1 = key + "_1"
+        val2 = key + "_2"
+        filter_cond += [f"{view_name}.{key} >= :{val1}", f"{view_name}.{key} <= :{val2}"]
+        values[val1] = float(distance_grid[0])
+        values[val2] = float(distance_grid[1])
 
     if building is not None:
         key = "building_count"
-        if "WHERE" in filter_cond:
-            filter_cond = filter_cond + f" AND {view_name}.{key} >= {building[0]} AND" \
-                                        f" {view_name}.{key} <= {building[1]}"
-        else:
-            filter_cond = f" WHERE {view_name}.{key} >= {building[0]} AND" \
-                          f" {view_name}.{key} <= {building[1]}"
+        val1 = key + "_1"
+        val2 = key + "_2"
+        filter_cond += [f"{view_name}.{key}>=:{val1}", f"{view_name}.{key}<=:{val2}"]
+        values[val1] = int(building[0])
+        values[val2] = int(building[1])
 
     if buildingfp is not None:
         key = "percentage_building_area"
-        if "WHERE" in filter_cond:
-            filter_cond = filter_cond + f" AND {view_name}.{key} >= {buildingfp[0]} AND" \
-                                        f" {view_name}.{key} <= {buildingfp[1]}"
-        else:
-            filter_cond = f" WHERE {view_name}.{key} >= {buildingfp[0]} AND" \
-                          f" {view_name}.{key} <= {buildingfp[1]}"
+        val1 = key + "_1"
+        val2 = key + "_2"
+        filter_cond += [f"{view_name}.{key}>=:{val1}", f"{view_name}.{key}<=:{val2}"]
+        values[val1] = float(buildingfp[0])
+        values[val2] = float(buildingfp[1])
 
     if keys is None:
         columns = "*"
@@ -130,9 +148,15 @@ def filter_materialized_view(
             columns = ", ".join(keys)
         else:
             columns = "COUNT({})".format(keys)
+
+    if len(filter_cond) > 0:
+        filter_cond_str = " WHERE " + " AND ".join(filter_cond)
+    else:
+        filter_cond_str = ""
+
     with engine.connect() as con:
-        query = 'SELECT {} FROM {}{}{};'.format(columns, view_name, filter_cond, limit)
-        rs = con.execute(query)
+        query = 'SELECT {} FROM {}{}{};'.format(columns, view_name, filter_cond_str, limit)
+        rs = con.execute(text(query), **values)
         data = rs.fetchall()
     return data