Initial depth check setup, testing has also been completed

api/app.py

@@ -3,6 +3,9 @@
 from flask import Flask
 from flask_graphql import GraphQLView
 from waitress import serve
+
+from backend.security_check import SecurityAnalysisBackend
+
 from db import (
     db,
     DB_NAME,
@@ -15,19 +18,21 @@
 
 app = Flask(__name__)
 
-app.config[
-    'SQLALCHEMY_DATABASE_URI'] = f'postgresql+psycopg2://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}'
+app.config['SQLALCHEMY_DATABASE_URI'] = f'postgresql+psycopg2://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}'
 app.config['SQLALCHEMY_COMMIT_ON_TEARDOWN'] = True
 app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True
 app.debug = True
 
 db.init_app(app)
 
+backend = SecurityAnalysisBackend(max_depth=10, max_cost=1000)
+
 app.add_url_rule(
     '/graphql',
     view_func=GraphQLView.as_view(
         'graphql',
         schema=schema,
+        backend=backend,
         graphiql=True
     )
 )

api/backend/__init__.py

@@ -0,0 +1,24 @@
+from typing import (
+    Dict,
+    List
+)
+from graphql.language.ast import (
+    FragmentDefinition,
+    OperationDefinition
+)
+
+
+def get_fragments(definitions) -> Dict[str, FragmentDefinition]:
+    return {
+        definition.name.value: definition
+        for definition in definitions
+        if isinstance(definition, FragmentDefinition)
+    }
+
+
+def get_queries_and_mutations(definitions) -> List[OperationDefinition]:
+    return [
+        definition
+        for definition in definitions
+        if isinstance(definition, OperationDefinition)
+    ]

api/backend/cost_check.py

@@ -0,0 +1,72 @@
+from typing import (
+    Dict
+)
+from graphql.language.ast import (
+    Document,
+    FragmentDefinition,
+    OperationDefinition,
+    Node,
+    FragmentSpread,
+    Field,
+    InlineFragment
+)
+from backend import (
+    get_fragments,
+    get_queries_and_mutations
+)
+
+from backend.cost_map import cost_map
+
+
+class CostLimitReached(Exception):
+    pass
+
+
+def measure_cost(node: Node, fragments: Dict[str, FragmentDefinition]) -> int:
+    """
+    A function which recursively measures the cost of a Graphene Query
+    :type node: Node
+    :param node: Graphql-core object used for query traversal/indexing
+    :type fragments: dict
+    :param fragments: The fragments of the query
+    :rtype: int
+    :return: The cost of the node
+    """
+    if isinstance(node, FragmentSpread):
+        fragment = fragments.get(node.name.value)
+        return measure_cost(node=fragment, fragments=fragments)
+
+    elif isinstance(node, Field):
+        if node.name.value.lower() in ["__schema", "__introspection"]:
+            return 0
+        if not node.selection_set:
+            return cost_map.get(node.name.value, 1)
+        costs = []
+        for selection in node.selection_set.selections:
+            cost = measure_cost(node=selection, fragments=fragments)
+            costs.append(cost)
+        return sum(costs) + cost_map.get(node.name.value, 1)
+    elif (
+        isinstance(node, FragmentDefinition)
+        or isinstance(node, OperationDefinition)
+        or isinstance(node, InlineFragment)
+    ):
+        costs = []
+        for selection in node.selection_set.selections:
+            cost = measure_cost(node=selection, fragments=fragments)
+            costs.append(cost)
+        return sum(costs)
+    else:
+        raise Exception("Unknown node")
+
+
+def check_cost_analysis(max_cost: int, document: Document):
+    fragments = get_fragments(document.definitions)
+    queries = get_queries_and_mutations(document.definitions)
+
+    for query in queries:
+        total_cost = measure_cost(query, fragments)
+        if total_cost > max_cost:
+            raise CostLimitReached(
+                'Query cost is too high'
+            )

api/backend/cost_map.py

@@ -0,0 +1,5 @@
+cost_map = {
+    'getSectorById': 1,
+    'getOrgById': 1,
+    'getGroupById': 1
+}

api/backend/depth_check.py

@@ -0,0 +1,69 @@
+from typing import Dict
+from graphql.language.ast import (
+    Document,
+    FragmentDefinition,
+    OperationDefinition,
+    Node,
+    FragmentSpread,
+    Field,
+    InlineFragment
+)
+
+from backend import (
+    get_fragments,
+    get_queries_and_mutations
+)
+
+
+class DepthLimitReached(Exception):
+    pass
+
+
+def measure_depth(node: Node, fragments: Dict[str, FragmentDefinition]) -> int:
+    """
+    A function which recursively measures the depth of a Graphene Query
+    :type node: Node
+    :param node: Graphql-core object used for query traversal/indexing
+    :type fragments: dict
+    :param fragments: The fragments of the query
+    :rtype: int
+    :return: The max depth of the node
+    """
+    if isinstance(node, FragmentSpread):
+        fragment = fragments.get(node.name.value)
+        return measure_depth(node=fragment, fragments=fragments)
+
+    elif isinstance(node, Field):
+        if node.name.value.lower() in ["__schema", "__introspection"]:
+            return 0
+        if not node.selection_set:
+            return 1
+        depths = []
+        for selection in node.selection_set.selections:
+            depth = measure_depth(node=selection, fragments=fragments)
+            depths.append(depth)
+        return 1 + max(depths)
+    elif (
+        isinstance(node, FragmentDefinition)
+        or isinstance(node, OperationDefinition)
+        or isinstance(node, InlineFragment)
+    ):
+        depths = []
+        for selection in node.selection_set.selections:
+            depth = measure_depth(node=selection, fragments=fragments)
+            depths.append(depth)
+        return max(depths)
+    else:
+        raise Exception("Unknown node")
+
+
+def check_max_depth(max_depth: int, document: Document):
+    fragments = get_fragments(document.definitions)
+    queries = get_queries_and_mutations(document.definitions)
+
+    for query in queries:
+        depth = measure_depth(query, fragments)
+        if depth > max_depth:
+            raise DepthLimitReached(
+                'Query is too complex'
+            )

api/backend/security_check.py

@@ -0,0 +1,24 @@
+from typing import (
+    Union,
+    Optional,
+    Any
+)
+from graphql import GraphQLDocument, GraphQLSchema
+from graphql.backend.core import GraphQLCoreBackend
+from graphql.language.ast import Document
+
+from backend.depth_check import check_max_depth
+from backend.cost_check import check_cost_analysis
+
+
+class SecurityAnalysisBackend(GraphQLCoreBackend):
+    def __init__(self, max_depth=10, max_cost=1000, executor: Optional[Any] = None):
+        super().__init__(executor=executor)
+        self.max_depth = max_depth
+        self.max_cost = max_cost
+
+    def document_from_string(self, schema: GraphQLSchema, document_string: Union[Document, str]) -> GraphQLDocument:
+        document = super().document_from_string(schema, document_string)
+        check_max_depth(max_depth=self.max_depth, document=document.document_ast)
+        check_cost_analysis(max_cost=self.max_cost, document=document.document_ast)
+        return document

api/resolvers/sectors.py

@@ -1,38 +1,41 @@
 from graphql import GraphQLError
 from schemas.sectors import Sectors, SectorsModel
 from model_enums.sectors import SectorEnums
+from manage import app
 
 
 # Resolvers
 def resolve_get_sector_by_id(self, info, **kwargs):
-	"""Return a sector by its row ID"""
-	sector_id = kwargs.get('id', 1)
-	query = Sectors.get_query(info).filter(
-		SectorsModel.id == sector_id
-	)
-	if not len(query.all()):
-		raise GraphQLError("Error, Invalid ID")
-	return query.all()
+    """Return a sector by its row ID"""
+    sector_id = kwargs.get('id', 1)
+    with app.app_context():
+        query = Sectors.get_query(info).filter(
+            SectorsModel.id == sector_id
+        )
+    if not len(query.all()):
+        raise GraphQLError("Error, Invalid ID")
+    return query.all()
 
 
 def resolve_get_sectors_by_sector(self, info, **kwargs):
-	"""Return a list of sectors by its sector"""
-	sector = kwargs.get('sector', 'EMPTY')
-	query = Sectors.get_query(info).filter(
-		SectorsModel.sector == sector
-	)
-
-	if not len(query.all()):
-		raise GraphQLError("Error, Sector does not exist")
-	return query.all()
+    """Return a list of sectors by its sector"""
+    sector = kwargs.get('sector', 'EMPTY')
+    with app.app_context():
+        query = Sectors.get_query(info).filter(
+            SectorsModel.sector == sector
+        )
+    if not len(query.all()):
+        raise GraphQLError("Error, Sector does not exist")
+    return query.all()
 
 
 def resolve_get_sector_by_zone(self, info, **kwargs):
-	"""Return a list of sectors by their zone"""
-	zone = kwargs.get('zone')
-	query = Sectors.get_query(info).filter(
-		SectorsModel.zone == zone
-	)
-	if not len(query.all()):
-		raise GraphQLError("Error, Zone does not exist")
-	return query.all()
+    """Return a list of sectors by their zone"""
+    zone = kwargs.get('zone')
+    with app.app_context():
+        query = Sectors.get_query(info).filter(
+            SectorsModel.zone == zone
+        )
+    if not len(query.all()):
+        raise GraphQLError("Error, Zone does not exist")
+    return query.all()