Merge pull request #201 from Bidaya0/feature/issue-199

fix: syntax
HXSecurity · Jan 7, 2022 · a1517d3 · a1517d3
2 parents 10612d9 + b79535e
commit a1517d3
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git a/apiserver/report/handler/saas_method_pool_handler.py b/apiserver/report/handler/saas_method_pool_handler.py
@@ -154,7 +154,7 @@ def save_method_call(self, pool_sign, current_version_agents):
             method_pool.res_header = utils.base64_decode(self.http_res_header)
             method_pool.res_body = decode_content(
                 get_res_body(self.http_res_body, self.version),
-                get_content_encoding(self.http_req_header))
+                get_content_encoding(self.http_req_header),self.version)
             method_pool.uri_sha1 = self.sha1(self.http_uri)
             method_pool.save(update_fields=[
                 'update_time',
@@ -193,7 +193,7 @@ def save_method_call(self, pool_sign, current_version_agents):
                 res_header=utils.base64_decode(self.http_res_header),
                 res_body = decode_content(
                 get_res_body(self.http_res_body, self.version),
-                get_content_encoding(self.http_req_header))
+                get_content_encoding(self.http_req_header),self.version),
                 context_path=self.context_path,
                 method_pool=json.dumps(self.method_pool),
                 pool_sign=pool_sign,
@@ -232,10 +232,12 @@ def sha1(raw):
         return h.hexdigest()
 
 
-def decode_content(body, content_type):
+def decode_content(body, content_type, version):
+    if version == 'v1':
+        return body
     if content_type == 'gzip':
         try:
-            return gzip.decompress(bytes(body, encoding='utf-8'))
+            return gzip.decompress(body)
         except:
             logger.error('not gzip type but using gzip as content_encoding')
     try:
@@ -259,8 +261,8 @@ def get_content_encoding(header):
 
 def get_res_body(res_body, version):
     if version == 'v1':
-        return res_body
+        return res_body #bytes
     elif version == 'v2':
-        return base64.b64decode(res_body)
+        return base64.b64decode(res_body) #bytes
     logger.info('no match version now version: {}'.format(version))
     return res_body
diff --git a/test/apiserver/test_agent_method_pool.py b/test/apiserver/test_agent_method_pool.py
@@ -11,7 +11,7 @@
 from dongtai.models.agent import IastAgent
 from dongtai.models.agent_method_pool import MethodPool
 import gzip
-
+import base64
 
 class AgentMethodPoolTestCase(AgentTestCase):
 
@@ -1039,8 +1039,8 @@ def test_agent_method_pool_gzip_test(self):
         testdata = '11231231321331232131231312233hwqeqqwe'
         data['detail'][
             'reqHeader'] = "Q29udGVudC1UeXBlPWFwcGxpY2F0aW9uL2pzb24KWC1GcmFtZS1PcHRpb25zPURFTlkKQ29udGVudC1MZW5ndGg9NjYKQ29udGVudC1lbmNvZGluZz1nemlwClgtQ29udGVudC1UeXBlLU9wdGlvbnM9bm9zbmlmZgpSZWZlcnJlci1Qb2xpY3k9c2FtZS1vcmlnaW4="
-        data['detail']['resBody'] = gzip_test_data = str(
-            gzip.compress(bytes(testdata, encoding='utf-8')))
+        data['version'] = 'v2'
+        data['detail']['resBody'] = gzip_test_data = base64.b64encode(gzip.compress(bytes(testdata, encoding='utf-8'))).decode('raw_unicode_escape')
         data = gzipdata(data)
         response = self.client.post(
             'http://testserver/api/v1/report/upload',
@@ -1052,11 +1052,16 @@ def test_agent_method_pool_gzip_test(self):
         assert MethodPool.objects.filter(
             url="http://localhost:9999/sqli123132123313132321123231test",
             agent_id=self.agent_id).exists()
-        assert MethodPool.objects.filter(
-            url="http://localhost:9999/sqli123132123313132321123231test",
-            agent_id=self.agent_id,
-            res_body=gzip_test_data).exists()
         assert not MethodPool.objects.filter(
             url="http://localhost:9999/sqli123132123313132321123231test",
             agent_id=self.agent_id,
-            res_body=testdata).exists()
+            res_body=gzip_test_data).exists()
+
+        #print(MethodPool.objects.filter(
+        #    url="http://localhost:9999/sqli123132123313132321123231test",
+        #    agent_id=self.agent_id,
+        #    res_body=testdata).values('res_body',flat=True).all()
+        #assert MethodPool.objects.filter(
+        #    url="http://localhost:9999/sqli123132123313132321123231test",
+        #    agent_id=self.agent_id,
+        #    res_body=testdata).exists()