From 3790c6c75617374ca094fc8f70bdc93024cf374e Mon Sep 17 00:00:00 2001
From: Jason Wen <haibin.wen3@gmail.com>
Date: Thu, 8 Jun 2023 08:48:37 -0400
Subject: [PATCH] Fleet Manager: One-Time Password cleanup (#177)

* init

* keep previous page from session
---
 system/fleetmanager/fleet_manager.py | 37 ++++++++++++----------------
 system/fleetmanager/helpers.py       | 12 +++++++++
 2 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/system/fleetmanager/fleet_manager.py b/system/fleetmanager/fleet_manager.py
index 0c245b0d0b8db2..c5e40f79b1181a 100755
--- a/system/fleetmanager/fleet_manager.py
+++ b/system/fleetmanager/fleet_manager.py
@@ -4,21 +4,11 @@
 import secrets
 import system.fleetmanager.helpers as fleet
 from flask import Flask, render_template, Response, request, send_from_directory, session, redirect, url_for
-from functools import wraps
 from system.loggerd.config import ROOT as REALDATA
 
 app = Flask(__name__)
 
 
-def login_required(f):
-  @wraps(f)
-  def decorated_route(*args, **kwargs):
-    if not session.get("logged_in"):
-      return redirect(url_for("index_page"))
-    return f(*args, **kwargs)
-  return decorated_route
-
-
 @app.route("/")
 def index_page():
   if session.get("logged_in"):
@@ -34,20 +24,24 @@ def login():
 
   if inputted_pin == correct_pin:
     session["logged_in"] = True
-    return redirect(url_for("home_page"))
+    if "previous_page" in session:
+      previous_page = session["previous_page"]
+      session.pop("previous_page", None)
+      return redirect(previous_page)
+    else:
+      return redirect(url_for("home_page"))
   else:
     error_message = "Incorrect PIN. Please try again."
     return render_template("login.html", error=error_message)
 
 
 @app.route("/index")
-@login_required
 def home_page():
   return render_template("index.html")
 
 
 @app.route("/footage/full/<cameratype>/<route>")
-@login_required
+@fleet.login_required
 def full(cameratype, route):
   chunk_size = 1024 * 512  # 5KiB
   file_name = cameratype + (".ts" if cameratype == "qcamera" else ".hevc")
@@ -61,7 +55,7 @@ def generate_buffered_stream():
 
 
 @app.route("/footage/<cameratype>/<segment>")
-@login_required
+@fleet.login_required
 def fcamera(cameratype, segment):
   if not fleet.is_valid_segment(segment):
     return render_template("error.html", error="invalid segment")
@@ -70,7 +64,7 @@ def fcamera(cameratype, segment):
 
 
 @app.route("/footage/<route>")
-@login_required
+@fleet.login_required
 def route(route):
   if len(route) != 20:
     return render_template("error.html", error="route not found")
@@ -90,14 +84,16 @@ def route(route):
   return render_template("route.html", route=route, query_type=query_type, links=links, segments=segments, query_segment=query_segment)
 
 
+@app.route("/footage/")
 @app.route("/footage")
-@login_required
+@fleet.login_required
 def footage():
   return render_template("footage.html", rows=fleet.all_routes())
 
 
+@app.route("/screenrecords/")
 @app.route("/screenrecords")
-@login_required
+@fleet.login_required
 def screenrecords():
   rows = fleet.all_screenrecords()
   if not rows:
@@ -106,26 +102,25 @@ def screenrecords():
 
 
 @app.route("/screenrecords/<clip>")
-@login_required
+@fleet.login_required
 def screenrecord(clip):
   return render_template("screenrecords.html", rows=fleet.all_screenrecords(), clip=clip)
 
 
 @app.route("/screenrecords/play/pipe/<file>")
-@login_required
+@fleet.login_required
 def videoscreenrecord(file):
   file_name = fleet.SCREENRECORD_PATH + file
   return Response(fleet.ffplay_mp4_wrap_process_builder(file_name).stdout.read(), status=200, mimetype='video/mp4')
 
 
 @app.route("/screenrecords/download/<clip>")
-@login_required
+@fleet.login_required
 def download_file(clip):
   return send_from_directory(fleet.SCREENRECORD_PATH, clip, as_attachment=True)
 
 
 @app.route("/about")
-@login_required
 def about():
   return render_template("about.html")
 
diff --git a/system/fleetmanager/helpers.py b/system/fleetmanager/helpers.py
index 604aebde466a3c..04191727198d55 100644
--- a/system/fleetmanager/helpers.py
+++ b/system/fleetmanager/helpers.py
@@ -1,5 +1,7 @@
 import os
 import subprocess
+from flask import render_template, request, session
+from functools import wraps
 from pathlib import Path
 from system.hardware import PC
 from system.loggerd.config import ROOT as REALDATA
@@ -15,6 +17,16 @@
   PIN_PATH = "/data/otp/"
 
 
+def login_required(f):
+  @wraps(f)
+  def decorated_route(*args, **kwargs):
+    if not session.get("logged_in"):
+      session["previous_page"] = request.url
+      return render_template("login.html")
+    return f(*args, **kwargs)
+  return decorated_route
+
+
 def all_screenrecords():
   return sorted(listdir_by_creation(SCREENRECORD_PATH), reverse=True)