From e9c64d6deb098a038eb641019bfe3eaa63b28c3e Mon Sep 17 00:00:00 2001
From: lash <dev@holbrook.no>
Date: Mon, 24 Jan 2022 19:46:25 +0000
Subject: [PATCH] Handle 404 error

---
 cic_auth_helper/error.py           |  7 +++++++
 cic_auth_helper/proxy.py           | 12 +++++++++++-
 cic_auth_helper/runnable/server.py | 13 +++++++++++--
 3 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100644 cic_auth_helper/error.py

diff --git a/cic_auth_helper/error.py b/cic_auth_helper/error.py
new file mode 100644
index 0000000..7e12009
--- /dev/null
+++ b/cic_auth_helper/error.py
@@ -0,0 +1,7 @@
+class NotFoundError(FileNotFoundError):
+
+    def __init__(self, code, msg, headers):
+        super(NotFoundError, self).__init__(msg)
+        self.code = code
+        self.msg = msg
+        self.headers = headers
diff --git a/cic_auth_helper/proxy.py b/cic_auth_helper/proxy.py
index 46f5534..93aa633 100644
--- a/cic_auth_helper/proxy.py
+++ b/cic_auth_helper/proxy.py
@@ -4,6 +4,9 @@ import urllib.request
 import os
 import logging
 
+# local imports
+from cic_auth_helper.error import NotFoundError
+
 logg = logging.getLogger(__name__)
 
 
@@ -15,6 +18,12 @@ def add_x_headers(env, header_f):
             header_orig = m[1].replace('_', '-')
             header_f(header_orig, env[x])
 
+
+class ReverseProxyHandler(urllib.request.BaseHandler):
+    
+    def http_error_404(self, request, response, code, msg, hdrs):
+        raise NotFoundError(code, msg, response.getheaders())
+
     
 class ReverseProxy:
 
@@ -25,6 +34,7 @@ class ReverseProxy:
         self.ignore_proxy_headers = []
         for h in ignore_proxy_headers:
             self.ignore_proxy_headers.append(h.lower())
+        self.opener = urllib.request.build_opener(ReverseProxyHandler())
 
 
     def proxy_pass(self, env, headers=[]):
@@ -34,7 +44,7 @@ class ReverseProxy:
         add_x_headers(env, req.add_header)
         req.add_header('Content-Type', env.get('CONTENT_TYPE', 'application/octet-stream'))
         req.data = env.get('wsgi.input')
-        res = urllib.request.urlopen(req)
+        res = self.opener.open(req)
 
         logg.debug('headers before reverse proxy {}'.format(headers))
 
diff --git a/cic_auth_helper/runnable/server.py b/cic_auth_helper/runnable/server.py
index 2686fdf..7994336 100644
--- a/cic_auth_helper/runnable/server.py
+++ b/cic_auth_helper/runnable/server.py
@@ -23,6 +23,9 @@ from usumbufu.retrieve.file import FileFetcher
 from usumbufu.adapters.uwsgi import UWSGIHTTPAuthorization
 from usumbufu.adapters.uwsgi import UWSGIAdapter
 
+# local imports
+from cic_auth_helper.error import NotFoundError
+
 logging.basicConfig(level=logging.WARNING)
 logg = logging.getLogger()
 
@@ -135,6 +138,7 @@ def application(env, start_response):
         #    return [b'']
         (challenge, expire) = challenger.request(env['REMOTE_ADDR'])
         headers.append(('WWW-Authenticate', hoba_auth_request_string(challenge, expire.timestamp(), realm=config.get('HTTP_AUTH_REALM'))),)
+        logg.debug('headers {}'.format(headers))
         start_response('401 authenticate or I will SCREAM_SNAKE_CASE at you', headers)
         return [b'']
 
@@ -156,9 +160,14 @@ def application(env, start_response):
     response_status = '200 OK'
     content = b''
     if reverse_proxy != None:
-        (response_status, headers, content) = reverse_proxy.proxy_pass(env, headers)
+        try:
+            (response_status, headers, content) = reverse_proxy.proxy_pass(env, headers)
+        except NotFoundError as e:
+            response_status = '404 ' + e.msg
+            headers = e.headers
+            content = b''
     else:
         content = str(auth_resource).encode('utf-8')
 
-    start_response('200 OK', headers)
+    start_response(response_status, headers)
     return [content]