From 13eed3fb33af3bf2a1cc6a203234bbd8953539ea Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Wed, 4 Dec 2024 22:57:41 +0100
Subject: [PATCH] fixup! feat(core): implement THP

---
 core/src/trezor/workflow.py | 37 ++++++++++++++++++++++---------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/core/src/trezor/workflow.py b/core/src/trezor/workflow.py
index 9fc72c3e98..a7a4791b1d 100644
--- a/core/src/trezor/workflow.py
+++ b/core/src/trezor/workflow.py
@@ -18,22 +18,29 @@ if __debug__:
     from trezor import utils
 
 if utils.USE_THP:
-    protocol_specific = ThpMessageType.ThpCreateNewSession
+    ALLOW_WHILE_LOCKED = (
+        ThpMessageType.ThpCreateNewSession,
+        MessageType.EndSession,
+        MessageType.GetFeatures,
+        MessageType.Cancel,
+        MessageType.LockDevice,
+        MessageType.DoPreauthorized,
+        MessageType.WipeDevice,
+        MessageType.SetBusy,
+        MessageType.Ping,
+    )
 else:
-    protocol_specific = MessageType.Initialize
-
-
-ALLOW_WHILE_LOCKED = (
-    protocol_specific,
-    MessageType.EndSession,
-    MessageType.GetFeatures,
-    MessageType.Cancel,
-    MessageType.LockDevice,
-    MessageType.DoPreauthorized,
-    MessageType.WipeDevice,
-    MessageType.SetBusy,
-    MessageType.Ping,
-)
+    ALLOW_WHILE_LOCKED = (
+        MessageType.Initialize,
+        MessageType.EndSession,
+        MessageType.GetFeatures,
+        MessageType.Cancel,
+        MessageType.LockDevice,
+        MessageType.DoPreauthorized,
+        MessageType.WipeDevice,
+        MessageType.SetBusy,
+        MessageType.Ping,
+    )
 
 
 # Set of workflow tasks.  Multiple workflows can be running at the same time.