From fd6313ea2a548d5642705c33f29f6a89eac97de9 Mon Sep 17 00:00:00 2001
From: Amir <ahosseinpour@asec.io>
Date: Mon, 24 Jun 2024 01:23:33 -0400
Subject: [PATCH] added GraphQLQuery class for detection purposes

---
 graphw00f/lib.py | 42 ++++++++++++++++++++++++++++--------------
 queries.json     | 14 ++++++++++++++
 2 files changed, 42 insertions(+), 14 deletions(-)
 create mode 100644 queries.json

diff --git a/graphw00f/lib.py b/graphw00f/lib.py
index a202094..bbf2199 100644
--- a/graphw00f/lib.py
+++ b/graphw00f/lib.py
@@ -1,4 +1,4 @@
-import requests
+import requests, json
 
 from graphw00f.helpers import error_contains
 
@@ -9,6 +9,21 @@ class GraphQLDetectionFailed(Exception):
 
 class GraphQLError(Exception):
   pass
+class GraphQLQuery:
+    def __init__(self, name, query, expected_response):
+        self.name = name
+        self.query = query
+        self.expected_response = expected_response
+
+    def __repr__(self):
+        return f"GraphQLQuery(name={self.name}, query={self.query}, expected_response={self.expected_response})"
+
+    @classmethod
+    def from_json(cls, json_file):
+        with open(json_file, 'r') as file:
+            data = json.load(file)
+        queries = data['queries']
+        return [cls(query['name'], query['query'], query['expected_response']) for query in queries]
 
 class GRAPHW00F:
   def __init__(self, headers,
@@ -24,20 +39,19 @@ def __init__(self, headers,
     self.proxies = proxies
 
   def check(self, url):
-    query = '''
-      query {
-        __typename
-      }
-    '''
-    response = self.graph_query(url, payload=query)
-
-    if response.get('data'):
-      if response.get('data', {}).get('__typename', '') in ('Query', 'QueryRoot', 'query_root'):
+    # Parse JSON file and create list of GraphQLQuery objects
+    query_objects = GraphQLQuery.from_json('queries.json')
+
+    # Check for each query in Query Objects
+    for query_object in query_objects:
+      response = self.graph_query(url, payload=query_object.query)
+      if response.get('data'):
+        if response.get('data', {}).get(query_object.expected_response, '') in ('Query', 'QueryRoot', 'query_root'):
+          return True
+      elif response.get('errors') and (any('locations' in i for i in response['errors']) or any('extensions' in i for i in response['errors'])):
+        return True
+      elif response.get('data'):
         return True
-    elif response.get('errors') and (any('locations' in i for i in response['errors']) or any('extensions' in i for i in response['errors'])):
-      return True
-    elif response.get('data'):
-      return True
     else:
       raise GraphQLDetectionFailed
 
diff --git a/queries.json b/queries.json
new file mode 100644
index 0000000..5c90993
--- /dev/null
+++ b/queries.json
@@ -0,0 +1,14 @@
+{
+    "queries": [
+        {
+            "name": "typenameQuery",
+            "query":"query{__typename}",
+            "expected_response":"__typename"
+        },
+        {
+            "name": "introspectionQuery",
+            "query": "query { __schema {types {name}}}",
+            "expected_response": "__schema"
+        }
+    ]
+}
\ No newline at end of file