From b493a3842c528d2b350deabfdbb5300435c19de2 Mon Sep 17 00:00:00 2001
From: Alexandre ACEBEDO <alexandre@acebedo.fr>
Date: Sun, 27 Sep 2020 14:36:34 +0200
Subject: [PATCH] Added command line options to change the rofi options
 depending on the window to open

---
 bitwarden_pyro/bwpyro.py         | 14 +++++++++++---
 bitwarden_pyro/util/arguments.py | 17 ++++++++++++++++-
 bitwarden_pyro/view/rofi.py      | 26 +++++++++++++++-----------
 3 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/bitwarden_pyro/bwpyro.py b/bitwarden_pyro/bwpyro.py
index a4ccb16..5355434 100644
--- a/bitwarden_pyro/bwpyro.py
+++ b/bitwarden_pyro/bwpyro.py
@@ -3,6 +3,8 @@
 import re
 import sys
 import logging
+import shlex
+from collections import namedtuple
 
 from bitwarden_pyro.util.logger import ProjectLogger
 from bitwarden_pyro.util.arguments import parse_arguments
@@ -72,7 +74,7 @@ def __lock(self):
             self._session.lock()
         except SessionException:
             self._logger.exception("Failed to lock session")
-            self._rofi = Rofi(None, None, None)
+            self._rofi = Rofi([], None, None)
             self._rofi.show_error("Failed to lock and delete session")
 
     def __unlock(self, force=False):
@@ -197,8 +199,14 @@ def __init_ui(self):
         try:
             self._config = ConfigLoader(self._args)
             self._session = Session(
-                self._config.get_int('security.timeout'))
-            self._rofi = Rofi(self._args.rofi_args,
+            self._config.get_int('security.timeout'))
+
+            RofiArgs = namedtuple('RofiArgs',
+                                  'main_window_args password_window_args additional_args')
+            rofi_args = RofiArgs(shlex.split(self._args.main_window_rofi_args),
+                                 shlex.split(self._args.password_window_rofi_args),
+                                 self._args.rofi_args)
+            self._rofi = Rofi(rofi_args,
                               self._config.get_itemaction('keyboard.enter'),
                               self._config.get_boolean('interface.hide_mesg'))
             self._clipboard = Clipboard(
diff --git a/bitwarden_pyro/util/arguments.py b/bitwarden_pyro/util/arguments.py
index 639a943..4e3b727 100644
--- a/bitwarden_pyro/util/arguments.py
+++ b/bitwarden_pyro/util/arguments.py
@@ -118,10 +118,25 @@ def parse_arguments():
         choices=['uris', 'logins', 'names', 'folders']
     )
 
+    parser.add_argument(
+        "--main-window-rofi-args",
+        help="rofi arguments used for the main window",
+        type=str,
+        default=""
+    )
+
+    parser.add_argument(
+        "--password-window-rofi-args",
+        help="rofi arguments used for the password window",
+        type=str,
+        default=""
+    )
+
     parser.add_argument(
         'rofi_args',
         help=argparse.SUPPRESS,
-        nargs=argparse.REMAINDER
+        nargs=argparse.REMAINDER,
+        default=[]
     )
 
     return parser.parse_args()
diff --git a/bitwarden_pyro/view/rofi.py b/bitwarden_pyro/view/rofi.py
index 45f2afd..841237c 100644
--- a/bitwarden_pyro/view/rofi.py
+++ b/bitwarden_pyro/view/rofi.py
@@ -14,17 +14,21 @@ class Rofi:
     def __init__(self, args, enter_event, hide_mesg):
         self._logger = ProjectLogger().get_logger()
         self._keybinds = {}
-        self._args = args[1:]
+        self._main_window_args = args.main_window_args + args.additional_args[1:]
+        self._password_window_args = args.password_window_args + args.additional_args[1:]
         self._enter_event = enter_event
         self._hide_mesg = hide_mesg
         self._keybinds_code = 10
 
-        if len(args) > 0:
-            self._logger.debug("Setting rofi arguments: %s", self._args)
+        if len(self._main_window_args) > 0:
+            self._logger.debug("Setting rofi arguments for main window: %s", self._main_window_args)
+        if len(self._password_window_args) > 0:
+            self._logger.debug("Setting rofi arguments for password window: %s",
+                                self._password_window_args)
 
-    def __extend_command(self, command):
-        if len(self._args) > 0:
-            command.extend(self._args)
+    def __extend_command(self, command, args):
+        if len(args) > 0:
+            command.extend(args)
 
         if not self._hide_mesg:
             mesg = []
@@ -68,8 +72,8 @@ def get_password(self):
                 "-password", "-lines", "0"
             ]
 
-            if len(self._args) > 0:
-                cmd.extend(self._args)
+            if len(self._password_window_args) > 0:
+                cmd.extend(self._password_window_args)
 
             proc = sp.run(cmd, check=True, capture_output=True)
             return proc.stdout.decode("utf-8").strip()
@@ -84,8 +88,8 @@ def show_error(self, message):
             self._logger.info("Showing Rofi error")
             cmd = ["rofi", "-e", f"ERROR! {message}"]
 
-            if len(self._args) > 0:
-                cmd.extend(self._args)
+            if len(self._main_window_args) > 0:
+                cmd.extend(self._main_window_args)
 
             sp.run(cmd, capture_output=True, check=True)
         except CalledProcessError:
@@ -99,7 +103,7 @@ def show_items(self, items, prompt='Bitwarden'):
             echo_cmd = ["echo", items]
             rofi_cmd = self.__extend_command([
                 "rofi", "-dmenu", "-p", prompt, "-i", "-no-custom"
-            ])
+            ], self._main_window_args)
 
             echo_proc = sp.Popen(echo_cmd, stdout=sp.PIPE)
             rofi_proc = sp.run(