#6607 - ckeygen should provide a default for the keyfile

src/twisted/conch/scripts/ckeygen.py

@@ -9,10 +9,13 @@
 
 import getpass
 import os
+import platform
 import socket
 import sys
+from collections.abc import Callable
 from functools import wraps
 from imp import reload
+from typing import Any, Dict, Optional
 
 from twisted.conch.ssh import keys
 from twisted.python import failure, filepath, log, usage
@@ -118,7 +121,7 @@ def enumrepresentation(options):
         return options
     else:
         raise keys.BadFingerPrintFormat(
-            "Unsupported fingerprint format: {}".format(options["format"])
+            f"Unsupported fingerprint format: {options['format']}"
         )
 
 
@@ -202,45 +205,71 @@ def _defaultPrivateKeySubtype(keyType):
         return "PEM"
 
 
-def printFingerprint(options):
-    if not options["filename"]:
-        filename = os.path.expanduser("~/.ssh/id_rsa")
-        options["filename"] = input("Enter file in which the key is (%s): " % filename)
-    if os.path.exists(options["filename"] + ".pub"):
-        options["filename"] += ".pub"
+def _getKeyOrDefault(
+    options: Dict[Any, Any],
+    inputCollector: Optional[Callable] = None,
+    keyTypeName: str = "rsa",
+) -> str:
+    """
+    If C{options["filename"]} is None, prompt the user to enter a path
+    or attempt to set it to .ssh/id_rsa
+    @param options: command line options
+    @param inputCollector: dependency injection for testing
+    @param keyTypeName: key type or "rsa"
+    """
+    if inputCollector is None:
+        inputCollector = input
+    filename = options["filename"]
+    if not filename:
+        filename = os.path.expanduser(f"~/.ssh/id_{keyTypeName}")
+        if platform.system() == "Windows":
+            filename = os.path.expanduser(fR"%HOMEPATH %\.ssh\id_{keyTypeName}")
+        filename = (
+            inputCollector("Enter file in which the key is (%s): " % filename)
+            or filename
+        )
+    return str(filename)
+
+
+def printFingerprint(options: Dict[Any, Any]) -> None:
+    filename = _getKeyOrDefault(options)
+    if os.path.exists(filename + ".pub"):
+        filename += ".pub"
     options = enumrepresentation(options)
     try:
-        key = keys.Key.fromFile(options["filename"])
+        key = keys.Key.fromFile(filename)
         print(
             "%s %s %s"
             % (
                 key.size(),
                 key.fingerprint(options["format"]),
-                os.path.basename(options["filename"]),
+                os.path.basename(filename),
             )
         )
     except keys.BadKeyError:
         sys.exit("bad key")
+    except FileNotFoundError:
+        sys.exit(f"{filename} could not be opened, please specify a file.")
 
 
 def changePassPhrase(options):
-    if not options["filename"]:
-        filename = os.path.expanduser("~/.ssh/id_rsa")
-        options["filename"] = input("Enter file in which the key is (%s): " % filename)
+    filename = _getKeyOrDefault(options)
     try:
-        key = keys.Key.fromFile(options["filename"])
+        key = keys.Key.fromFile(filename)
     except keys.EncryptedKeyError:
         # Raised if password not supplied for an encrypted key
         if not options.get("pass"):
             options["pass"] = getpass.getpass("Enter old passphrase: ")
         try:
-            key = keys.Key.fromFile(options["filename"], passphrase=options["pass"])
+            key = keys.Key.fromFile(filename, passphrase=options["pass"])
         except keys.BadKeyError:
             sys.exit("Could not change passphrase: old passphrase error")
         except keys.EncryptedKeyError as e:
             sys.exit(f"Could not change passphrase: {e}")
     except keys.BadKeyError as e:
         sys.exit(f"Could not change passphrase: {e}")
+    except FileNotFoundError:
+        sys.exit(f"{filename} could not be opened, please specify a file.")
 
     if not options.get("newpass"):
         while 1:
@@ -268,22 +297,22 @@ def changePassPhrase(options):
     except (keys.EncryptedKeyError, keys.BadKeyError) as e:
         sys.exit(f"Could not change passphrase: {e}")
 
-    with open(options["filename"], "wb") as fd:
+    with open(filename, "wb") as fd:
         fd.write(newkeydata)
 
     print("Your identification has been saved with the new passphrase.")
 
 
 def displayPublicKey(options):
-    if not options["filename"]:
-        filename = os.path.expanduser("~/.ssh/id_rsa")
-        options["filename"] = input("Enter file in which the key is (%s): " % filename)
+    filename = _getKeyOrDefault(options)
     try:
-        key = keys.Key.fromFile(options["filename"])
+        key = keys.Key.fromFile(filename)
+    except FileNotFoundError:
+        sys.exit(f"{filename} could not be opened, please specify a file.")
     except keys.EncryptedKeyError:
         if not options.get("pass"):
             options["pass"] = getpass.getpass("Enter passphrase: ")
-        key = keys.Key.fromFile(options["filename"], passphrase=options["pass"])
+        key = keys.Key.fromFile(filename, passphrase=options["pass"])
     displayKey = key.public().toString("openssh").decode("ascii")
     print(displayKey)
 
@@ -297,29 +326,36 @@ def _inputSaveFile(prompt: str) -> str:
     return input(prompt)
 
 
-def _saveKey(key, options):
+def _saveKey(
+    key: keys.Key,
+    options: Dict[Any, Any],
+    inputCollector: Optional[Callable] = None,
+) -> None:
     """
     Persist a SSH key on local filesystem.
 
     @param key: Key which is persisted on local filesystem.
-    @type key: C{keys.Key} implementation.
 
     @param options:
-    @type options: L{dict}
+
+    @param inputCollector: Dependency injection for testing.
     """
+    if inputCollector is None:
+        inputCollector = input
     KeyTypeMapping = {"EC": "ecdsa", "Ed25519": "ed25519", "RSA": "rsa", "DSA": "dsa"}
     keyTypeName = KeyTypeMapping[key.type()]
-    if not options["filename"]:
-        defaultPath = os.path.expanduser(f"~/.ssh/id_{keyTypeName}")
+    filename = options["filename"]
+    if not filename:
+        defaultPath = _getKeyOrDefault(options, inputCollector, keyTypeName)
         newPath = _inputSaveFile(
             f"Enter file in which to save the key ({defaultPath}): "
         )
 
-        options["filename"] = newPath.strip() or defaultPath
+        filename = newPath.strip() or defaultPath
 
-    if os.path.exists(options["filename"]):
-        print("{} already exists.".format(options["filename"]))
-        yn = input("Overwrite (y/n)? ")
+    if os.path.exists(filename):
+        print(f"{filename} already exists.")
+        yn = inputCollector("Overwrite (y/n)? ")
         if yn[0].lower() != "y":
             sys.exit()
 
@@ -339,23 +375,24 @@ def _saveKey(key, options):
 
     comment = f"{getpass.getuser()}@{socket.gethostname()}"
 
-    filepath.FilePath(options["filename"]).setContent(
+    fp = filepath.FilePath(filename)
+    fp.setContent(
         key.toString(
             "openssh",
             subtype=options["private-key-subtype"],
             passphrase=options["pass"],
         )
     )
-    os.chmod(options["filename"], 33152)
+    fp.chmod(0o100600)
 
-    filepath.FilePath(options["filename"] + ".pub").setContent(
+    filepath.FilePath(filename + ".pub").setContent(
         key.public().toString("openssh", comment=comment)
     )
     options = enumrepresentation(options)
 
-    print("Your identification has been saved in {}".format(options["filename"]))
-    print("Your public key has been saved in {}.pub".format(options["filename"]))
-    print("The key fingerprint in {} is:".format(options["format"]))
+    print(f"Your identification has been saved in {filename}")
+    print(f"Your public key has been saved in {filename}.pub")
+    print(f"The key fingerprint in {options['format']} is:")
     print(key.fingerprint(options["format"]))
 
 

src/twisted/conch/test/test_ckeygen.py

@@ -6,6 +6,7 @@
 """
 
 import getpass
+import os
 import subprocess
 import sys
 from io import StringIO
@@ -23,6 +24,7 @@
 
 if requireModule("cryptography") and requireModule("pyasn1"):
     from twisted.conch.scripts.ckeygen import (
+        _getKeyOrDefault,
         _saveKey,
         changePassPhrase,
         displayPublicKey,
@@ -172,6 +174,19 @@ def test_printFingerprintBadFingerPrintFormat(self):
             "Unsupported fingerprint format: sha-base64", em.exception.args[0]
         )
 
+    def test_printFingerprintSuffixAppended(self) -> None:
+        """
+        L{printFingerprint} checks if the filename with the  '.pub' suffix
+        exists in ~/.ssh.
+        """
+        filename = self.mktemp()
+        FilePath(filename + ".pub").setContent(publicRSA_openssh)
+        printFingerprint({"filename": filename, "format": "md5-hex"})
+        self.assertEqual(
+            self.stdout.getvalue(),
+            "2048 85:25:04:32:58:55:96:9f:57:ee:fb:a8:1a:ea:69:da temp.pub\n",
+        )
+
     def test_saveKey(self):
         """
         L{_saveKey} writes the private and public parts of a key to two
@@ -345,17 +360,43 @@ def test_saveKeyNoFilename(self):
         base = FilePath(self.mktemp())
         base.makedirs()
         keyPath = base.child("custom_key").path
+        input_prompts = []
 
         import twisted.conch.scripts.ckeygen
 
+        def mock_input(*args):
+            return input_prompts.append("")
+
         self.patch(twisted.conch.scripts.ckeygen, "_inputSaveFile", lambda _: keyPath)
         key = Key.fromString(privateRSA_openssh)
-        _saveKey(key, {"filename": None, "no-passphrase": True, "format": "md5-hex"})
+        _saveKey(
+            key,
+            {"filename": None, "no-passphrase": True, "format": "md5-hex"},
+            mock_input,
+        )
 
         persistedKeyContent = base.child("custom_key").getContent()
         persistedKey = key.fromString(persistedKeyContent, None, b"")
         self.assertEqual(key, persistedKey)
 
+    def test_saveKeyFileExists(self) -> None:
+        """
+        When the specified file exists, it will ask the user for confirmation
+        before overwriting.
+        """
+
+        def mock_input(*args):
+            return ["n"]
+
+        base = FilePath(self.mktemp())
+        base.makedirs()
+        keyPath = base.child("custom_key").path
+
+        self.patch(os.path, "exists", lambda _: True)
+        key = Key.fromString(privateRSA_openssh)
+        options = {"filename": keyPath, "no-passphrase": True, "format": "md5-hex"}
+        self.assertRaises(SystemExit, _saveKey, key, options, mock_input)
+
     def test_saveKeySubtypeV1(self):
         """
         L{_saveKey} can be told to write the new private key file in OpenSSH
@@ -628,3 +669,52 @@ def test_changePassphraseSubtypeV1(self):
         self.assertTrue(
             privateKeyContent.startswith(b"-----BEGIN OPENSSH PRIVATE KEY-----\n")
         )
+
+    def test_useDefaultForKey(self) -> None:
+        """
+        L{options} will default to "~/.ssh/id_rsa" if the user doesn't
+        specify a key.
+        """
+        input_prompts = []
+
+        def mock_input(*args):
+            return input_prompts.append("")
+
+        options = {"filename": ""}
+
+        filename = _getKeyOrDefault(options, mock_input)
+        self.assertEqual(
+            options["filename"],
+            "",
+        )
+        # Resolved path is an RSA key inside .ssh dir.
+        self.assertTrue(filename.endswith(os.path.join(".ssh", "id_rsa")))
+        # The user is prompted once to enter the path, since no path was
+        # provided via CLI.
+        self.assertEqual(1, len(input_prompts))
+        self.assertEqual([""], input_prompts)
+
+    def test_displayPublicKeyHandleFileNotFound(self) -> None:
+        """
+        Ensure FileNotFoundError is handled, whether the user has supplied
+        a bad path, or has no key at the default path.
+        """
+        options = {"filename": "/foo/bar"}
+        exc = self.assertRaises(SystemExit, displayPublicKey, options)
+        self.assertIn("could not be opened, please specify a file.", exc.args[0])
+
+    def test_changePassPhraseHandleFileNotFound(self) -> None:
+        """
+        Ensure FileNotFoundError is handled for an invalid filename.
+        """
+        options = {"filename": "/foo/bar"}
+        exc = self.assertRaises(SystemExit, changePassPhrase, options)
+        self.assertIn("could not be opened, please specify a file.", exc.args[0])
+
+    def test_printFingerprintHandleFileNotFound(self) -> None:
+        """
+        Ensure FileNotFoundError is handled for an invalid filename.
+        """
+        options = {"filename": "/foo/bar", "format": "md5-hex"}
+        exc = self.assertRaises(SystemExit, printFingerprint, options)
+        self.assertIn("could not be opened, please specify a file.", exc.args[0])

src/twisted/newsfragments/6607.bugfix

@@ -0,0 +1 @@
+twisted.conch.scripts.ckeygen now substitutes a default of "~/.ssh/id_rsa" if no keyfile is specified.