From 80e2cdcd8d83f71f34fbfc96776f744c41303f95 Mon Sep 17 00:00:00 2001
From: Simone Margaritelli <evilsocket@gmail.com>
Date: Tue, 5 Nov 2019 14:26:35 +0100
Subject: [PATCH] refact: using flask templating
---
pwnagotchi/ui/display.py | 5 +-
pwnagotchi/ui/web/__init__.py | 12 +++
pwnagotchi/ui/{web.py => web/handler.py} | 95 +++++++-----------------
pwnagotchi/ui/web/server.py | 43 +++++++++++
4 files changed, 83 insertions(+), 72 deletions(-)
create mode 100644 pwnagotchi/ui/web/__init__.py
rename pwnagotchi/ui/{web.py => web/handler.py} (63%)
create mode 100644 pwnagotchi/ui/web/server.py
diff --git a/pwnagotchi/ui/display.py b/pwnagotchi/ui/display.py
index 7d077a0..8abb1d7 100644
--- a/pwnagotchi/ui/display.py
+++ b/pwnagotchi/ui/display.py
@@ -5,6 +5,7 @@ import threading
import pwnagotchi.plugins as plugins
import pwnagotchi.ui.hw as hw
import pwnagotchi.ui.web as web
+from pwnagotchi.ui.web.server import Server
from pwnagotchi.ui.view import View
@@ -15,7 +16,7 @@ class Display(View):
self._enabled = config['enabled']
self._rotation = config['rotation']
- self._webui = web.Server(config)
+ self._webui = Server(config)
self.init_display()
@@ -41,7 +42,7 @@ class Display(View):
def is_waveshare27inch(self):
return self._implementation.name == 'waveshare27inch'
-
+
def is_waveshare29inch(self):
return self._implementation.name == 'waveshare29inch'
diff --git a/pwnagotchi/ui/web/__init__.py b/pwnagotchi/ui/web/__init__.py
new file mode 100644
index 0000000..fd43206
--- /dev/null
+++ b/pwnagotchi/ui/web/__init__.py
@@ -0,0 +1,12 @@
+from threading import Lock
+
+frame_path = '/root/pwnagotchi.png'
+frame_format = 'PNG'
+frame_ctype = 'image/png'
+frame_lock = Lock()
+
+
+def update_frame(img):
+ global frame_lock, frame_path, frame_format
+ with frame_lock:
+ img.save(frame_path, format=frame_format)
diff --git a/pwnagotchi/ui/web.py b/pwnagotchi/ui/web/handler.py
similarity index 63%
rename from pwnagotchi/ui/web.py
rename to pwnagotchi/ui/web/handler.py
index c7396ff..a04cc54 100644
--- a/pwnagotchi/ui/web.py
+++ b/pwnagotchi/ui/web/handler.py
@@ -1,6 +1,3 @@
-import _thread
-import secrets
-from threading import Lock
import logging
import os
@@ -9,27 +6,14 @@ logging.getLogger('werkzeug').setLevel(logging.ERROR)
os.environ['WERKZEUG_RUN_MAIN'] = 'true'
import pwnagotchi
+import pwnagotchi.ui.web as web
from pwnagotchi.agent import Agent
from pwnagotchi import plugins
-from flask import Flask
+
from flask import send_file
from flask import request
from flask import abort
from flask import render_template_string
-from flask_cors import CORS
-from flask_wtf.csrf import CSRFProtect
-
-frame_path = '/root/pwnagotchi.png'
-frame_format = 'PNG'
-frame_ctype = 'image/png'
-frame_lock = Lock()
-
-
-def update_frame(img):
- global frame_lock, frame_path, frame_format
- with frame_lock:
- img.save(frame_path, format=frame_format)
-
STYLE = """
.block {
@@ -59,26 +43,27 @@ window.onload = function() {
function updateImage() {
image.src = image.src.split("?")[0] + "?" + new Date().getTime();
}
- setInterval(updateImage, %d);
+ setInterval(updateImage, 1000);
}
"""
INDEX = """<html>
<head>
- <title>%s</title>
+ <title>{{ title }}</title>
<style>""" + STYLE + """</style>
</head>
<body>
- <div style="position: absolute; top:0; left:0; width:100%%;" class="pixelated">
- <img src="/ui" id="ui" style="width:100%%;"/>
+ <div style="position: absolute; top:0; left:0; width:100%;" class="pixelated">
+ <img src="/ui" id="ui" style="width:100%;"/>
<br/>
<hr/>
<form style="display:inline;" method="POST" action="/shutdown" onsubmit="return confirm('This will halt the unit, continue?');">
<input style="display:inline;" type="submit" class="block" value="Shutdown"/>
<input type="hidden" name="csrf_token" value="{{ csrf_token() }}"/>
</form>
- <form style="display:inline;" method="POST" action="/restart" onsubmit="return confirm('This will restart the service in %s mode, continue?');">
- <input style="display:inline;" type="submit" class="block" value="Restart in %s mode"/>
+ <form style="display:inline;" method="POST" action="/restart" onsubmit="return confirm('This will restart the service in {{ other_mode }} mode, continue?');">
+ <input style="display:inline;" type="submit" class="block" value="Restart in {{ other_mode }} mode"/>
+ <input type="hidden" name="mode" value="{{ other_mode }}"/>
<input type="hidden" name="csrf_token" value="{{ csrf_token() }}"/>
</form>
</div>
@@ -89,18 +74,18 @@ INDEX = """<html>
STATUS_PAGE = """<html>
<head>
- <title>%s</title>
+ <title>{{ title }}</title>
<style>""" + STYLE + """</style>
</head>
<body>
- <div style="position: absolute; top:0; left:0; width:100%%;">
- %s
+ <div style="position: absolute; top:0; left:0; width:100%;">
+ {{ message }}
</div>
</body>
</html>"""
-class RequestHandler:
+class Handler:
def __init__(self, app):
self._app = app
self._app.add_url_rule('/', 'index', self.index)
@@ -115,12 +100,8 @@ class RequestHandler:
self._app.add_url_rule('/plugins/<name>/<path:subpath>', 'plugins', self.plugins, methods=['GET', 'POST'])
def index(self):
- other_mode = 'AUTO' if Agent.INSTANCE.mode == 'manual' else 'MANU'
- return render_template_string(INDEX % (
- pwnagotchi.name(),
- other_mode,
- other_mode,
- 1000))
+ return render_template_string(INDEX, title=pwnagotchi.name(),
+ other_mode='AUTO' if Agent.INSTANCE.mode == 'manual' else 'MANU')
def plugins(self, name, subpath):
if name is None:
@@ -142,49 +123,23 @@ class RequestHandler:
# serve a message and shuts down the unit
def shutdown(self):
try:
- return render_template_string(STATUS_PAGE % (pwnagotchi.name(), 'Shutting down ...'))
+ return render_template_string(STATUS_PAGE, title=pwnagotchi.name(), message='Shutting down ...')
finally:
pwnagotchi.shutdown()
# serve a message and restart the unit in the other mode
def restart(self):
- other_mode = 'AUTO' if Agent.INSTANCE.mode == 'manual' else 'MANU'
+ mode = request.form['mode']
+ if mode not in ('AUTO', 'MANU'):
+ mode = 'MANU'
+
try:
- return render_template_string(STATUS_PAGE % (pwnagotchi.name(), 'Restart in %s mode ...' % other_mode))
+ return render_template_string(STATUS_PAGE, title=pwnagotchi.name(),
+ message='Restart in %s mode ...' % mode)
finally:
- pwnagotchi.restart(other_mode)
+ pwnagotchi.restart(mode)
# serve the PNG file with the display image
def ui(self):
- global frame_lock, frame_path
- with frame_lock:
- return send_file(frame_path, mimetype='image/png')
-
-
-class Server:
- def __init__(self, config):
- self._enabled = config['video']['enabled']
- self._port = config['video']['port']
- self._address = config['video']['address']
- self._origin = None
-
- if 'origin' in config['video']:
- self._origin = config['video']['origin']
-
- if self._enabled:
- _thread.start_new_thread(self._http_serve, ())
-
- def _http_serve(self):
- if self._address is not None:
- app = Flask(__name__)
- app.secret_key = secrets.token_urlsafe(256)
-
- if self._origin:
- CORS(app, resources={r"*": {"origins": self._origin}})
-
- CSRFProtect(app)
- RequestHandler(app)
-
- app.run(host=self._address, port=self._port, debug=False)
- else:
- logging.info("could not get ip of usb0, video server not starting")
+ with web.frame_lock:
+ return send_file(web.frame_path, mimetype='image/png')
diff --git a/pwnagotchi/ui/web/server.py b/pwnagotchi/ui/web/server.py
new file mode 100644
index 0000000..5bc1fef
--- /dev/null
+++ b/pwnagotchi/ui/web/server.py
@@ -0,0 +1,43 @@
+import _thread
+import secrets
+import logging
+import os
+
+# https://stackoverflow.com/questions/14888799/disable-console-messages-in-flask-server
+logging.getLogger('werkzeug').setLevel(logging.ERROR)
+os.environ['WERKZEUG_RUN_MAIN'] = 'true'
+
+from flask import Flask
+from flask_cors import CORS
+from flask_wtf.csrf import CSRFProtect
+
+from pwnagotchi.ui.web.handler import Handler
+
+
+class Server:
+ def __init__(self, config):
+ self._enabled = config['video']['enabled']
+ self._port = config['video']['port']
+ self._address = config['video']['address']
+ self._origin = None
+
+ if 'origin' in config['video']:
+ self._origin = config['video']['origin']
+
+ if self._enabled:
+ _thread.start_new_thread(self._http_serve, ())
+
+ def _http_serve(self):
+ if self._address is not None:
+ app = Flask(__name__)
+ app.secret_key = secrets.token_urlsafe(256)
+
+ if self._origin:
+ CORS(app, resources={r"*": {"origins": self._origin}})
+
+ CSRFProtect(app)
+ Handler(app)
+
+ app.run(host=self._address, port=self._port, debug=False)
+ else:
+ logging.info("could not get ip of usb0, video server not starting")