Skip to content

Commit 70a80ec

Browse files
codeskyblueCopilot
andauthored
支持 SSH 密钥密码短语验证,添加 cryptography 依赖 (#2)
* 支持 SSH 密钥密码短语验证,添加 cryptography 依赖 - 新增 _resolve_key_passphrase() 函数,通过 cryptography 库解析密钥并提示输入密码短语 - HostConfig 添加 get_password() 方法处理整数类型密码 - 添加 poetry.toml 配置 in-project 虚拟环境 - 更新 dev-dependencies 为新的 poetry group 格式 * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
1 parent b23d73e commit 70a80ec

3 files changed

Lines changed: 50 additions & 3 deletions

File tree

poetry.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
[virtualenvs]
2+
in-project = true

pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,12 @@ dataclasses-json = "*"
1818
pyyaml = "*"
1919
prompt_toolkit = "*"
2020
pexpect = "*"
21+
cryptography = "<43"
2122

2223
[tool.poetry.scripts]
2324
sshg = "sshg:main"
2425

25-
[tool.poetry.dev-dependencies]
26+
[tool.poetry.group.dev.dependencies]
2627
pytest = "^7.2.0"
2728
pytest-cov = "^2"
2829

sshg.py

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,11 @@ class HostConfig(DataClassJsonMixin):
9696
via: typing.Optional["HostConfig"] = make_field(mm_field=fields.Field(), default=None)
9797
_parent: typing.Optional["HostConfig"] = make_field(mm_field=fields.Field(), default=None, init=False, repr=False)
9898

99+
def get_password(self) -> str:
100+
if isinstance(self.password, int):
101+
return str(self.password)
102+
return self.password or ""
103+
99104
def post_load(self):
100105
if self._parent:
101106
if not self.user:
@@ -151,6 +156,43 @@ def output_filter(line):
151156
s.interact()
152157

153158

159+
def _resolve_key_passphrase(keypath: pathlib.Path) -> str:
160+
"""Try loading key without passphrase, prompt until correct if needed."""
161+
from cryptography.hazmat.primitives.serialization import load_pem_private_key, load_ssh_private_key
162+
163+
key_data = keypath.read_bytes()
164+
165+
def _try_load(password: typing.Optional[bytes]) -> typing.Optional[bool]:
166+
"""Return True if the key loads, False if a passphrase is required/incorrect, None if the key is invalid."""
167+
errors: typing.List[Exception] = []
168+
for loader in (load_ssh_private_key, load_pem_private_key):
169+
try:
170+
loader(key_data, password=password)
171+
return True
172+
except (TypeError, ValueError) as e:
173+
errors.append(e)
174+
175+
msg = " ".join(str(e).lower() for e in errors)
176+
if any(k in msg for k in ("password", "passphrase", "bad decrypt", "incorrect")):
177+
return False
178+
return None
179+
180+
res = _try_load(None)
181+
if res is None:
182+
raise ValueError(f"Unsupported or invalid private key: {keypath}")
183+
if res:
184+
return ""
185+
186+
while True:
187+
password = getpass.getpass(f"Enter passphrase for key {keypath}: ")
188+
res = _try_load(password.encode())
189+
if res is None:
190+
raise ValueError(f"Unsupported or invalid private key: {keypath}")
191+
if res:
192+
return password
193+
print("Wrong passphrase, try again.")
194+
195+
154196
def spawn_ssh(host_config: HostConfig, is_local: bool = True, ssh_client: pxssh.pxssh = None, reset_prompt: bool = None) -> pxssh.pxssh:
155197
# https://pexpect.readthedocs.io/en/stable/api/pxssh.html
156198
cmdargs = host_config.build_cmdargs()
@@ -164,6 +206,8 @@ def spawn_ssh(host_config: HostConfig, is_local: bool = True, ssh_client: pxssh.
164206
if keypath.stat().st_mode & 0o077 != 0:
165207
print("Warning: keypath mode change to 0600")
166208
keypath.chmod(0o600)
209+
if not host_config.get_password():
210+
host_config.password = _resolve_key_passphrase(keypath)
167211

168212
s.SSH_OPTS += " -o StrictHostKeyChecking=no"
169213
if reset_prompt is None:
@@ -172,7 +216,7 @@ def spawn_ssh(host_config: HostConfig, is_local: bool = True, ssh_client: pxssh.
172216
if is_local:
173217
s.login(host_config.host,
174218
username=host_config.user,
175-
password=host_config.password,
219+
password=host_config.get_password(),
176220
port=host_config.port,
177221
ssh_key=keypath,
178222
quiet=False,
@@ -182,7 +226,7 @@ def spawn_ssh(host_config: HostConfig, is_local: bool = True, ssh_client: pxssh.
182226
else:
183227
s.login(host_config.host,
184228
username=host_config.user,
185-
password=host_config.password,
229+
password=host_config.get_password(),
186230
port=host_config.port,
187231
ssh_key=keypath,
188232
quiet=False,

0 commit comments

Comments
 (0)