add more http protocol processing APIs

gor/base.py

@@ -119,6 +119,26 @@ def http_status(self, payload):
     def set_http_status(self, payload, new_status):
         return self.set_http_path(payload, new_status)
 
+    def http_headers(self, payload):
+        """
+        Parse the payload and return http headers in a map
+        :param payload: the http payload to inspect
+        :return: a map mapping from key to value of each http header item
+        """
+        if isinstance(payload, bytes):
+            payload = payload.decode()
+        start_index = payload.index("\r\n")
+        end_index = payload.index("\r\n\r\n")
+        if end_index == -1:
+            end_index = len(payload)
+        headers = {}
+        for item in payload[start_index+2:end_index].split("\r\n"):
+            sep_index = item.index(":")
+            key = item[:sep_index]
+            value = item[sep_index+1:].lstrip()
+            headers[key] = value
+        return headers
+
     def http_header(self, payload, name):
         current_line = 0
         idx = 0
@@ -168,6 +188,15 @@ def set_http_header(self, payload, name, value):
         else:
             return payload[:header['value_start']] + ' ' + value + '\r\n' + payload[header['end']:]
 
+    def delete_http_header(self, payload, name):
+        if isinstance(payload, bytes):
+            payload = payload.decode()
+        header = self.http_header(payload, name)
+        if header is None:
+            return payload
+        else:
+            return payload[:header['start']] + payload[header['end']:]
+
     def http_body(self, payload):
         if '\r\n\r\n' not in payload:
             return ''
@@ -194,3 +223,9 @@ def set_http_cookie(self, payload, name, value):
         cookies = list(filter(lambda x: not x.startswith(name + '='), cookies.split('; ')))
         cookies.append(name + '=' + value)
         return self.set_http_header(payload, 'Cookie', '; '.join(cookies))
+
+    def delete_http_cookie(self, payload, name):
+        cookie_data = self.http_header(payload, 'Cookie')
+        cookies = cookie_data.get('value') or ''
+        cookies = list(filter(lambda x: not x.startswith(name + '='), cookies.split('; ')))
+        return self.set_http_header(payload, 'Cookie', '; '.join(cookies))

tests/test_gor.py

@@ -54,6 +54,17 @@ def test_http_path_param(self):
         self.assertEqual(self.gor.http_path(payload), '/?test=123&qwer=ty')
         self.assertIn('ty', self.gor.http_path_param(payload, 'qwer'))
 
+    def test_http_headers(self):
+        payload = b'GET / HTTP/1.1\r\nHost: localhost:3000\r\nUser-Agent: Python\r\nContent-Length:5\r\n\r\nhello'
+        expected = {
+            'Host': 'localhost:3000',
+            'User-Agent': 'Python',
+            'Content-Length': '5',
+        }
+        headers = self.gor.http_headers(payload)
+        for k, v in expected.items():
+            self.assertEqual(headers.get(k), v)
+
     def test_http_header(self):
         payload = b'GET / HTTP/1.1\r\nHost: localhost:3000\r\nUser-Agent: Python\r\nContent-Length:5\r\n\r\nhello'
         expected = {
@@ -82,21 +93,34 @@ def test_set_http_header(self):
         new_payload = self.gor.set_http_header(new_payload, 'X-Test2', 'test2')
         self.assertEqual(new_payload, expected)
 
+    def test_delete_http_header(self):
+        payload = b'GET / HTTP/1.1\r\nUser-Agent: Python\r\nContent-Length: 5\r\n\r\nhello'
+        new_payload = self.gor.delete_http_header(payload, 'User-Agent')
+        self.assertEqual(new_payload, 'GET / HTTP/1.1\r\nContent-Length: 5\r\n\r\nhello')
+        new_payload = self.gor.delete_http_header(new_payload, 'not-exists-header')
+        self.assertEqual(new_payload, 'GET / HTTP/1.1\r\nContent-Length: 5\r\n\r\nhello')
+
     def test_http_body(self):
         payload = 'GET / HTTP/1.1\r\nUser-Agent: Python\r\nContent-Length: 5\r\n\r\nhello'
         body = self.gor.http_body(payload)
         self.assertEqual(body, 'hello')
 
+        invalid_payload = 'GET / HTTP/1.1\r\nUser-Agent: Python\r\nContent-Length: 5\r\nhello'
+        body = self.gor.http_body(invalid_payload)
+        self.assertEqual(body, '')
+
     def test_set_http_body(self):
         payload = 'GET / HTTP/1.1\r\nUser-Agent: Python\r\nContent-Length: 5\r\n\r\nhello'
         new_payload = self.gor.set_http_body(payload, 'hello, world!')
         expected = 'GET / HTTP/1.1\r\nUser-Agent: Python\r\nContent-Length: 13\r\n\r\nhello, world!'
         self.assertEqual(new_payload, expected)
 
     def test_http_cookie(self):
-        payload = 'GET / HTTP/1.1\r\nCookie: a=b; test=zxc\r\n\r\n'
+        payload = 'GET / HTTP/1.1\r\nCookie: a=b; test=zxc; test2=a=b\r\n\r\n'
         cookie = self.gor.http_cookie(payload, 'test')
         self.assertEqual(cookie, 'zxc')
+        cookie = self.gor.http_cookie(payload, 'test2')
+        self.assertEqual(cookie, 'a=b')
         cookie = self.gor.http_cookie(payload, 'nope')
         self.assertIsNone(cookie)
 
@@ -106,3 +130,8 @@ def test_set_http_cookie(self):
         self.assertEqual(new_payload, 'GET / HTTP/1.1\r\nCookie: a=b; test=111\r\n\r\n')
         new_payload = self.gor.set_http_cookie(payload, 'new', 'one%3d%3d--test')
         self.assertEqual(new_payload, 'GET / HTTP/1.1\r\nCookie: a=b; test=zxc; new=one%3d%3d--test\r\n\r\n')
+
+    def test_delete_http_cookie(self):
+        payload = b'GET / HTTP/1.1\r\nCookie: a=b; test=zxc\r\n\r\n'
+        new_payload = self.gor.delete_http_cookie(payload, 'test')
+        self.assertEqual(new_payload, 'GET / HTTP/1.1\r\nCookie: a=b\r\n\r\n')