refactor: Remove redundant code in run.py.

Author: Gopinath-Mahendiran

zulipterminal/cli/run.py

@@ -4,6 +4,7 @@
 
 import argparse
 import configparser
+import contextlib
 import logging
 import os
 import stat
@@ -311,67 +312,60 @@ def fetch_zuliprc(zuliprc_path: str) -> None:
         print(in_color("red", "\nIncorrect Email(or Username) or Password!\n"))
         login_data = get_api_key(realm_url)
 
+    zulip_key_path = os.path.join(
+        os.path.dirname(os.path.abspath(zuliprc_path)), "zulip_key"
+    )
+
     preferred_realm_url, login_id, api_key = login_data
     save_zuliprc_failure = _write_zuliprc(
-        zuliprc_path,
+        to_path=zuliprc_path,
+        key_path=zulip_key_path,
+        api_key=api_key,
         login_id=login_id,
         server_url=preferred_realm_url,
     )
-    zulip_key_path = os.path.join(
-        os.path.dirname(os.path.abspath(zuliprc_path)), "zulip_key"
-    )
-    save_zulipkey_failure = _write_zulip_key(zulip_key_path, api_key=api_key)
-    if not save_zulipkey_failure:
-        print(f"Generated API key saved at {zulip_key_path}")
-    else:
-        exit_with_error(save_zuliprc_failure)
-
     if not save_zuliprc_failure:
         print(f"Generated config file saved at {zuliprc_path}")
     else:
         exit_with_error(save_zuliprc_failure)
 
 
-def _write_zuliprc(to_path: str, *, login_id: str, server_url: str) -> str:
+def _write_zuliprc(
+    to_path: str, *, key_path: str, login_id: str, api_key: str, server_url: str
+) -> str:
     """
-    Writes a zuliprc file, returning a non-empty error string on failure
-    Only creates new private files; errors if file already exists
+    Writes both zuliprc and zulip_key files securely.
+    Ensures atomicity: if one file fails to write, cleans up the other.
+    Returns an empty string on success, or a descriptive error message on failure.
     """
+    zuliprc_created = False
+
     try:
+        # Write zuliprc
         with open(
             os.open(to_path, os.O_CREAT | os.O_WRONLY | os.O_EXCL, 0o600), "w"
         ) as f:
             f.write(
-                f"[api]\nemail={login_id}\n"
-                f"# Fill the passcmd field with a command that outputs the API key.\n"
-                f"# The API key is temporarily stored in the 'zulip_key' file at "
-                f"{os.path.dirname(os.path.abspath(to_path))}."
-                f"\n# After storing the key in a password manager, replace the cmd.\n"
-                f"passcmd=cat zulip_key\nsite={server_url}"
+                f"[api]\nemail={login_id}\npasscmd=cat zulip_key\nsite={server_url}"
             )
-        return ""
-    except FileExistsError:
-        return f"zuliprc already exists at {to_path}"
-    except OSError as ex:
-        return f"{ex.__class__.__name__}: zuliprc could not be created at {to_path}"
-
-
-def _write_zulip_key(to_path: str, *, api_key: str) -> str:
-    """
-    Writes a zulip_key file, returning a non-empty error string on failure
-    Only creates new private files; errors if file already exists
-    """
-
-    try:
+        zuliprc_created = True
+        # Write zulip_key
         with open(
-            os.open(to_path, os.O_CREAT | os.O_WRONLY | os.O_EXCL, 0o600), "w"
+            os.open(key_path, os.O_CREAT | os.O_WRONLY | os.O_EXCL, 0o600), "w"
         ) as f:
             f.write(api_key)
+
         return ""
+
     except FileExistsError:
-        return f"zulip_key already exists at {to_path}"
+        filename = to_path if not zuliprc_created else key_path
+        return f"FileExistsError: {filename} already exists"
     except OSError as ex:
-        return f"{ex.__class__.__name__}: zulip_key could not be created at {to_path}"
+        if zuliprc_created:
+            with contextlib.suppress(Exception):
+                os.remove(to_path)
+        filename = key_path if zuliprc_created else to_path
+        return f"{ex.__class__.__name__}: could not create {filename} ({ex})"
 
 
 def parse_zuliprc(zuliprc_str: str) -> Dict[str, SettingData]: