From a553df0d963793a694e14b4c069e78eae5f1f701 Mon Sep 17 00:00:00 2001
From: Sergey Kolupaev <sergey@keepersecurity.com>
Date: Tue, 18 Jun 2024 23:01:42 -0700
Subject: [PATCH] Custom UI sample

---
 examples/custom_login_ui.py | 42 +++++++++++++++++++++++++++++++++++++
 keepercommander/api.py      |  6 +++---
 2 files changed, 45 insertions(+), 3 deletions(-)
 create mode 100644 examples/custom_login_ui.py

diff --git a/examples/custom_login_ui.py b/examples/custom_login_ui.py
new file mode 100644
index 000000000..62abbd50c
--- /dev/null
+++ b/examples/custom_login_ui.py
@@ -0,0 +1,42 @@
+import logging
+
+from keepercommander.auth import login_steps
+from keepercommander import api
+from keepercommander.__main__ import get_params_from_config
+
+
+class CustomUi(login_steps.LoginUi):
+    def on_device_approval(self, step):
+        try:
+            step.send_push(login_steps.DeviceApprovalChannel.Email)
+        except Exception as e:
+            logging.error(e)
+        verification_code = wait_for_verification_code()
+        step.send_code(login_steps.DeviceApprovalChannel.Email, verification_code)
+
+    def on_two_factor(self, step):
+        raise NotImplementedError()
+
+    def on_password(self, step):
+        password = wait_for_password()
+        try:
+            step.verify_password(password)
+        except Exception as e:
+            logging.error(e)
+
+    def on_sso_redirect(self, step):
+        raise NotImplementedError()
+
+    def on_sso_data_key(self, step):
+        raise NotImplementedError()
+
+
+my_params = get_params_from_config('config.json')
+while not my_params.user:
+    my_params.user = input('User(Email): ')
+
+ui = CustomUi()
+api.login(my_params, login_ui=ui)
+
+if not my_params.session_token:
+    exit(1)
\ No newline at end of file
diff --git a/keepercommander/api.py b/keepercommander/api.py
index c8a1fb77d..008fbf491 100644
--- a/keepercommander/api.py
+++ b/keepercommander/api.py
@@ -60,11 +60,11 @@ def run_command(params, request):
     return rest_api.v2_execute(params.rest_context, request)
 
 
-def login(params, new_login=False):
-    # type: (KeeperParams, bool) -> None
+def login(params, new_login=False, login_ui=None):
+    # type: (KeeperParams, bool, Optional[Any]) -> None
 
     logging.info('Logging in to Keeper Commander')
-    flow = loginv3.LoginV3Flow()
+    flow = loginv3.LoginV3Flow(login_ui)
     try:
         flow.login(params, new_login=new_login)
     except loginv3.InvalidDeviceToken: