Repository: amgedr/webchk

Issue:
Add option to specify a header for all requests
I would like to be able to specify an authentication header for all of the URLs in my list. Could this feature be added? Thanks!

Comments:

Patch:
diff --git a/setup.py b/setup.py
index 2dd5cd5..456ffbf 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@ setup(
     version='1.1.0',
     packages=['webchk'],
     test_suite='test',
-    url='https://codehill.com/projects/webchk',
+    url='https://webchk.codehill.com',
     license="MIT license",
     author='Amged Rustom',
     author_email='amgadhs@codehill.com',
diff --git a/webchk/__init__.py b/webchk/__init__.py
index 12f478c..9997150 100644
--- a/webchk/__init__.py
+++ b/webchk/__init__.py
@@ -2,4 +2,4 @@ __version__ = '1.1.0'
 __cmd_description__ = """Check HTTP status codes, response headers and redirects.
 This is free software and it comes with absolutely no warranty.
 You can distribute and modify it under terms of MIT License.
-Homepage: https://codehill.com/projects/webchk"""
+Homepage: https://webchk.codehill.com"""
diff --git a/webchk/__main__.py b/webchk/__main__.py
index ff2671f..f35ff14 100644
--- a/webchk/__main__.py
+++ b/webchk/__main__.py
@@ -1,20 +1,18 @@
 import sys
 import threading
 
-from .utils import get_parser, read_input_file
+from .utils import get_parser, read_input_file, format_headers
 from .http import http_response, HTTPRequests
 from . import __version__
 
 
-def _process_url(url, requests, get_request):
-    resp = http_response(
-        url=url,
-        timeout=requests.timeout,
-        parse=requests.parse_xml,
-        get_request=get_request,
-    )
+def _process_url(url, requests: HTTPRequests):
+    resp = http_response(url, requests)
     print(resp, file=requests.output_file)
 
+    if requests.show_headers:
+        print('{}\n'.format(format_headers(resp.headers)))
+
     follow = resp.redirect
     while follow:
         print('   {}'.format(follow), file=requests.output_file)
@@ -44,11 +42,8 @@ def process_urls(requests: HTTPRequests):
             continue
 
         thread = threading.Thread(
-            target=_process_url, args=(
-                url,
-                requests,
-                requests.get_request,
-            )
+            target=_process_url,
+            args=(url, requests)
         )
         thread.start()
         threads.append(thread)
@@ -77,7 +72,10 @@ def main():
             list_only=args.list,
             parse_xml=args.parse,
             timeout=args.timeout,
+            show_headers=args.all,
             get_request=args.get,
+            user_agent=args.agent,
+            auth=args.auth,
         )
 
         if args.urls:
diff --git a/webchk/http.py b/webchk/http.py
index b6e0528..85491ea 100644
--- a/webchk/http.py
+++ b/webchk/http.py
@@ -1,20 +1,33 @@
-import collections
 import http.client
 from urllib.parse import urlparse
 import socket
 import ssl
+import sys
 import timeit
 
+from . import __version__
 from webchk.utils import urls_from_xml
 
 
-HTTPRequests = collections.namedtuple(
-    'HTTPRequests',
-    [
-        'urls', 'output_file', 'list_only', 'parse_xml', 'timeout',
-        'get_request',
-    ]
-)
+class HTTPRequests:
+    def __init__(self, urls, output_file=sys.stdout, list_only=False,
+                 parse_xml=False, timeout=3, show_headers=False,
+                 headers=None, get_request=False, auth=None,
+                 user_agent=None) -> None:
+        self.urls = urls
+        self.output_file = output_file
+        self.list_only = list_only
+        self.parse_xml = parse_xml
+        self.timeout = timeout
+        self.get_request = get_request
+        self.show_headers = show_headers
+
+        self.headers = headers if headers is not None else {}
+        if not user_agent:
+            user_agent = f'webchk v{__version__}'
+        self.headers['User-Agent'] = user_agent
+        if auth:
+            self.headers['Authorization'] = auth
 
 
 class Result:
@@ -79,16 +92,14 @@ def _http_connect(loc, timeout):
     return http.client.HTTPConnection(loc.netloc, timeout=timeout)
 
 
-def _http_request(loc, timeout, get_request=False):
-    """Performs a HTTP request and return response in a Result object.
-
-    Does a HEAD HTTP request if get_request is False and GET if True.
-    """
+def _http_request(loc, req: HTTPRequests):
+    """Performs a HTTP request and return response in a Result object."""
+    conn = None
     try:
-        conn = _http_connect(loc, timeout)
-        method = 'GET' if get_request else 'HEAD'
+        conn = _http_connect(loc, req.timeout)
+        method = 'GET' if req.get_request or req.parse_xml else 'HEAD'
 
-        conn.request(method, loc.path)
+        conn.request(method, loc.path, headers=req.headers)
         resp = conn.getresponse()
 
         result = Result(loc.geturl())
@@ -97,17 +108,19 @@ def _http_request(loc, timeout, get_request=False):
         result.fill_headers(resp.getheaders())
 
         # status code is not 204 (no content) and not a redirect
-        if get_request and resp.status not in (204, 301, 302, 303, 307, 308):
-            result.content = resp.read().decode('utf-8')
+        is_not_redirect = resp.status not in (204, 301, 302, 303, 307, 308)
+        if (req.get_request or req.parse_xml) and is_not_redirect:
+            result.content = resp.read()
 
     except TimeoutError:
         raise
     finally:
-        conn.close()
+        if conn:
+            conn.close()
     return result
 
 
-def http_response(url, timeout, parse=False, get_request=False):
+def http_response(url, requests: HTTPRequests):
     """Returns the HTTP response code.
 
     If the response code is a temporary or permanent redirect then it
@@ -125,10 +138,7 @@ def http_response(url, timeout, parse=False, get_request=False):
     try:
         start = timeit.default_timer()
 
-        # true if user wants HTTP GET or asked for the content to be parsed
-        force_get = get_request or (parse and url.endswith('.xml'))
-
-        result = _http_request(loc, timeout, get_request=force_get)
+        result = _http_request(loc, requests)
         result.latency = '{:2.3}'.format(timeit.default_timer() - start)
 
         if 400 <= result.status < 500:
@@ -152,25 +162,26 @@ def http_response(url, timeout, parse=False, get_request=False):
                 if new_url.startswith('/'):
                     new_url = '{}://{}{}'.format(
                         loc.scheme, loc.netloc, new_url)
-                result.redirect = http_response(
-                    new_url, timeout, parse=parse, get_request=get_request)
+                result.redirect = http_response(new_url, requests)
 
-        if result.content and parse:
+        if result.content and requests.parse_xml:
+            requests.parse_xml = False
             sitemap = urls_from_xml(result.content)
             result.sitemap_urls = []
             for s_url in sitemap:
                 # some sites include the sitemap's url in the sitemap
                 if s_url == result.url:
                     continue
-                result.sitemap_urls.append(
-                    http_response(s_url, timeout, get_request=get_request))
+                result.sitemap_urls.append(http_response(s_url, requests))
 
     except socket.gaierror:
         result.desc = 'Could not resolve'
     except (TimeoutError, socket.timeout):
         result.desc = 'Operation timed out'
-    except (http.client.RemoteDisconnected) as exc:
+    except http.client.RemoteDisconnected as exc:
         result.desc = str(exc)
+    except http.client.InvalidURL:
+        result.desc = 'Invalid URL'
     except (ConnectionRefusedError, ConnectionResetError) as exc:
         result.desc = exc.strerror
     except ssl.SSLCertVerificationError as exc:
diff --git a/webchk/utils.py b/webchk/utils.py
index e8508cf..2343445 100644
--- a/webchk/utils.py
+++ b/webchk/utils.py
@@ -29,6 +29,18 @@ def get_parser():
         help='Perform HTTP GET request instead of HEAD',
         action='store_true',
     )
+    parser.add_argument(
+        '--auth',
+        help='Set Authentication header',
+        type=str,
+        default='',
+    )
+    parser.add_argument(
+        '--agent',
+        help='Set a custom user-agent',
+        type=str,
+        default='',
+    )
     parser.add_argument(
         '-l', '--list',
         help='Print URLs without checking them',
@@ -77,3 +89,14 @@ def urls_from_xml(data):
             if j.tag.endswith("loc"):
                 urls.append(j.text.strip())
     return urls
+
+
+def format_headers(headers: dict) -> str:
+    if not isinstance(headers, dict):
+        raise ValueError
+
+    indent = ' '
+    formatted = []
+    for key, val in headers.items():
+        formatted.append('{}{}: {}'.format(indent, key, val))
+    return '\n'.join(formatted)


Test patch:
diff --git a/test/test_http.py b/test/test_http.py
index af6ac8e..f1a16da 100644
--- a/test/test_http.py
+++ b/test/test_http.py
@@ -1,6 +1,6 @@
 import unittest
 
-from webchk.http import http_response, parse_url
+from webchk.http import http_response, parse_url, HTTPRequests
 
 TIMEOUT = 3
 
@@ -26,12 +26,15 @@ class Http(unittest.TestCase):
 
     def test_http_response(self):
         for url, result in self.urls.items():
-            resp_code = http_response(url, TIMEOUT).status
+            req = HTTPRequests(url, timeout=TIMEOUT)
+            resp_code = http_response(url, req).status
             self.assertEqual(resp_code, result[0], url)
 
     def test_redirect_follows(self):
         url = 'https://httpstat.us/307'
-        resp = http_response(url, TIMEOUT)
+
+        req = HTTPRequests(url, timeout=TIMEOUT)
+        resp = http_response(url, req)
         total = 0
         while resp.redirect:
             fmt = '{} ... {} {} ({})'.format(
@@ -42,10 +45,12 @@ class Http(unittest.TestCase):
         self.assertEqual(total, 1)
 
     def test_unresolvable_domains(self):
-        resp = http_response('http://!.c', TIMEOUT)
+        req = HTTPRequests(None, timeout=TIMEOUT)
+        resp = http_response('http://!.c', req)
         self.assertEqual(str(resp), 'http://!.c ... Could not resolve')
 
     def test_timeouts(self):
         url = 'http://httpbin.org/delay/5'
-        resp = http_response(url, timeout=1)
+        req = HTTPRequests(url, timeout=1)
+        resp = http_response(url, req)
         self.assertEqual(resp.desc, 'Operation timed out')
diff --git a/test/test_utils.py b/test/test_utils.py
index 496022e..2ca9591 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -1,6 +1,8 @@
 import unittest
 
-from webchk.utils import read_input_file, get_parser, urls_from_xml
+from webchk.utils import (
+    read_input_file, get_parser, urls_from_xml, format_headers
+)
 
 
 class FileIOTest(unittest.TestCase):
@@ -46,3 +48,29 @@ class XmlParserTest(unittest.TestCase):
         ]
         urls = urls_from_xml(self.xml)
         self.assertEqual(urls, urls_list)
+
+
+class HeadersFormatterTest(unittest.TestCase):
+    def test_valid_headers(self):
+        cases = {
+            ' Connection: keep-alive\n Content-Length: 5386':
+            {
+                'Connection': 'keep-alive',
+                'Content-Length': '5386',
+            },
+            ' Cache-Control: no-cache\n Content-Type: text/html':
+            {
+                'Cache-Control': 'no-cache',
+                'Content-Type': 'text/html',
+            }
+        }
+
+        for expected, case in cases.items():
+            self.assertEqual(format_headers(case), expected)
+
+    def test_invalid_value(self):
+        cases = [[], 123, 'abc']
+
+        for case in cases:
+            with self.assertRaises(ValueError):
+                format_headers(case)