From 5643f9ae70ea51e38bb6afe63ac0c29872058435 Mon Sep 17 00:00:00 2001
From: python273 <iam@python273.pw>
Date: Sat, 26 Oct 2019 17:07:47 +0300
Subject: [PATCH] Set CORS headers only if set in config

---
 pwnagotchi/defaults.yml |  2 +-
 pwnagotchi/ui/web.py    | 33 +++++++++++++++++----------------
 pwnagotchi/utils.py     |  6 +++---
 3 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/pwnagotchi/defaults.yml b/pwnagotchi/defaults.yml
index e08906be3..ca0c74ca1 100644
--- a/pwnagotchi/defaults.yml
+++ b/pwnagotchi/defaults.yml
@@ -211,7 +211,7 @@ ui:
         video:
             enabled: true
             address: '0.0.0.0'
-            origin: '*'
+            origin: null
             port: 8080
             # command to be executed when a new png frame is available
             # for instance, to use with framebuffer based displays:
diff --git a/pwnagotchi/ui/web.py b/pwnagotchi/ui/web.py
index 8a142301f..716a2c0b7 100644
--- a/pwnagotchi/ui/web.py
+++ b/pwnagotchi/ui/web.py
@@ -75,7 +75,7 @@ def update_frame(img):
 
 
 class Handler(BaseHTTPRequestHandler):
-    AllowedOrigin = '*'
+    AllowedOrigin = None  # CORS headers are not sent
 
     # suppress internal logging
     def log_message(self, format, *args):
@@ -88,12 +88,13 @@ def _send_cors_headers(self):
         self.send_header("X-XSS-Protection", "1; mode=block")
         self.send_header("Referrer-Policy", "same-origin")
         # cors
-        self.send_header("Access-Control-Allow-Origin", Handler.AllowedOrigin)
-        self.send_header('Access-Control-Allow-Credentials', 'true')
-        self.send_header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
-        self.send_header("Access-Control-Allow-Headers",
-                         "Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization")
-        self.send_header("Vary", "Origin")
+        if Handler.AllowedOrigin:
+            self.send_header("Access-Control-Allow-Origin", Handler.AllowedOrigin)
+            self.send_header('Access-Control-Allow-Credentials', 'true')
+            self.send_header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
+            self.send_header("Access-Control-Allow-Headers",
+                            "Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization")
+            self.send_header("Vary", "Origin")
 
     # just render some html in a 200 response
     def _html(self, html):
@@ -132,15 +133,18 @@ def _image(self):
 
     # check the Origin header vs CORS
     def _is_allowed(self):
+        if not Handler.AllowedOrigin or Handler.AllowedOrigin == '*':
+            return True
+
+        # TODO: FIX doesn't work with GET requests same-origin
         origin = self.headers.get('origin')
-        if not origin and Handler.AllowedOrigin != '*':
+        if not origin:
             logging.warning("request with no Origin header from %s" % self.address_string())
             return False
 
-        if Handler.AllowedOrigin != '*':
-            if origin != Handler.AllowedOrigin:
-                logging.warning("request with blocked Origin from %s: %s" % (self.address_string(), origin))
-                return False
+        if origin != Handler.AllowedOrigin:
+            logging.warning("request with blocked Origin from %s: %s" % (self.address_string(), origin))
+            return False
 
         return True
 
@@ -186,11 +190,8 @@ def __init__(self, config):
         self._address = config['video']['address']
         self._httpd = None
 
-        if 'origin' in config['video'] and config['video']['origin'] != '*':
+        if 'origin' in config['video']:
             Handler.AllowedOrigin = config['video']['origin']
-        else:
-            logging.warning("THE WEB UI IS RUNNING WITH ALLOWED ORIGIN SET TO *, READ WHY YOU SHOULD CHANGE IT HERE " +
-                            "https://developer.mozilla.org/it/docs/Web/HTTP/CORS")
 
         if self._enabled:
             _thread.start_new_thread(self._http_serve, ())
diff --git a/pwnagotchi/utils.py b/pwnagotchi/utils.py
index 392d2fd7e..25aab13aa 100644
--- a/pwnagotchi/utils.py
+++ b/pwnagotchi/utils.py
@@ -79,7 +79,7 @@ def load_config(args):
     elif config['ui']['display']['type'] in ('papirus', 'papi'):
         config['ui']['display']['type'] = 'papirus'
 
-    elif config['ui']['display']['type'] in ('oledhat'):
+    elif config['ui']['display']['type'] in ('oledhat',):
         config['ui']['display']['type'] = 'oledhat'
 
     elif config['ui']['display']['type'] in ('ws_1', 'ws1', 'waveshare_1', 'waveshare1'):
@@ -91,9 +91,9 @@ def load_config(args):
     elif config['ui']['display']['type'] in ('ws_27inch', 'ws27inch', 'waveshare_27inch', 'waveshare27inch'):
         config['ui']['display']['type'] = 'waveshare27inch'
 
-    elif config['ui']['display']['type'] in ('lcdhat'):
+    elif config['ui']['display']['type'] in ('lcdhat',):
         config['ui']['display']['type'] = 'lcdhat'
-        
+
     else:
         print("unsupported display type %s" % config['ui']['display']['type'])
         exit(1)