Merge pull request #11 from Buzz2d0/fix#9

Fix#10

Author: boy-hack

HackRequests/HackRequests.py

@@ -149,6 +149,7 @@ def httpraw(self, raw: str, **kwargs):
         proxy = kwargs.get("proxy", None)
         real_host = kwargs.get("real_host", None)
         ssl = kwargs.get("ssl", False)
+        location = kwargs.get("location", True)
 
         scheme = 'http'
         port = 80
@@ -245,6 +246,13 @@ def httpraw(self, raw: str, **kwargs):
             _url = "{scheme}://{host}{path}".format(scheme=scheme, host=host, path=path)
         else:
             _url = "{scheme}://{host}{path}".format(scheme=scheme, host=host + ":" + port, path=path)
+        
+        redirect = rep.msg.get('location', None)  # handle 301/302
+        if redirect and location:
+            if not redirect.startswith('http'):
+                redirect = parse.urljoin(_url, redirect)
+            return self.http(redirect, post=None, method=method, headers=headers, location=True, locationcount=1)
+
         return response(rep, _url, log, )
 
     def http(self, url, **kwargs):
@@ -368,7 +376,7 @@ def __init__(self, rep, redirect, log, oldcookie=''):
             self.cookies = {}
 
         self.headers = _header_dict
-        self.header = self.rep.msg  # response header
+        self.header = str(self.rep.msg)  # response header
         self.log = log
         charset = self.rep.msg.get('content-type', 'utf-8')
         try:

tests/test_http_raw.py

@@ -89,3 +89,23 @@ def test_chunked(self):
         r = self.hack.httpraw(raw)
         print(r.text())
         self.assertTrue("helloworld!helloworld hellowor ld" in r.text())
+
+    def test_redirect(self):
+        raw = ''' 
+GET / HTTP/1.1
+Host: www.python.org
+Connection: Keep-Alive
+Cache-Control: max-age=0
+Upgrade-Insecure-Requests: 1
+User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/68.0.3440.106 Safari/537.36
+Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8
+Accept-Encoding: gzip, deflate
+Accept-Language: zh-CN,zh;q=0.9,en;q=0.8
+        '''
+        r = self.hack.httpraw(raw)
+        self.assertEqual(r.status_code, 200)
+        self.assertIn('class="python home"', r.text())
+
+        r = self.hack.httpraw(raw, location=False)
+        self.assertEqual(r.status_code, 301)
+        self.assertTrue(r.text() == "")