22
33import os
44import time
5+ import hmac
56import hashlib
67import asyncio
78import argparse
89import traceback
910from pathlib import Path
10- from typing import IO , Callable , cast , ContextManager
11+ from dataclasses import dataclass
12+ from typing import IO , Callable , TypeAlias , cast , ContextManager
1113from enum import Enum , auto
1214from concurrent .futures import ThreadPoolExecutor , CancelledError
1315
@@ -249,63 +251,112 @@ def ssh_terminal_handler(
249251 pass
250252
251253
254+ @dataclass
255+ class PasswordAndPublicKeyAuthentication :
256+ password : str
257+
258+
259+ @dataclass
260+ class PublicKeyAuthentication :
261+ pass
262+
263+
264+ @dataclass
265+ class NoAuthentication :
266+ pass
267+
268+
269+ AuthenticationMethod : TypeAlias = (
270+ PasswordAndPublicKeyAuthentication | PublicKeyAuthentication | NoAuthentication
271+ )
272+
273+
252274class SSHServer (asyncssh .SSHServer ):
253275 def __init__ (
254276 self ,
255- password : str | None ,
277+ authentication : AuthenticationMethod ,
256278 console_cls : type [Console ],
257279 namespace : argparse .Namespace ,
258280 executor : ThreadPoolExecutor ,
259281 ):
260282 self ._gambaterm_console_cls = console_cls
261283 self ._gambaterm_namespace = namespace
262284 self ._gambaterm_executor = executor
263- self ._gambaterm_password = password
285+ self ._gambaterm_authentication = authentication
264286
265287 def connection_made (self , conn : asyncssh .SSHServerConnection ) -> None :
266288 conn .set_extra_info (console_cls = self ._gambaterm_console_cls )
267289 conn .set_extra_info (executor = self ._gambaterm_executor )
268290 conn .set_extra_info (namespace = self ._gambaterm_namespace )
269291
270292 def begin_auth (self , username : str ) -> bool :
271- return True
293+ return not isinstance ( self . _gambaterm_authentication , NoAuthentication )
272294
273295 def session_requested (self ) -> SSHServerProcess [str ]:
274296 return asyncssh .SSHServerProcess (
275297 safe_ssh_process_handler , sftp_factory = None , sftp_version = 3 , allow_scp = False
276298 )
277299
278300 def password_auth_supported (self ) -> bool :
279- # Allow empty string as a valid password (--password ''),
280- # only None (unset) disables password auth.
281- return self . _gambaterm_password is not None
301+ return isinstance (
302+ self . _gambaterm_authentication , ( PasswordAndPublicKeyAuthentication ,)
303+ )
282304
283305 def validate_password (self , username : str , password : str ) -> bool :
284- assert self ._gambaterm_password is not None
285- return password == self ._gambaterm_password
306+ assert isinstance (
307+ self ._gambaterm_authentication , PasswordAndPublicKeyAuthentication
308+ )
309+ return hmac .compare_digest (password , self ._gambaterm_authentication .password )
286310
287311
288312async def run_server (
289313 bind : str ,
290314 port : int ,
291- password : str | None ,
315+ authentication : AuthenticationMethod ,
292316 console_cls : type [Console ],
293317 namespace : argparse .Namespace ,
294318 executor : ThreadPoolExecutor ,
295319) -> None :
296- ssh_key_dir = Path (os .environ .get ("GAMBATERM_SSH_KEY_DIR" , "~/.ssh" ))
297- user_private_key = (ssh_key_dir / "id_rsa" ).expanduser ()
298- user_public_key = (ssh_key_dir / "id_rsa.pub" ).expanduser ()
299- if not user_private_key .exists ():
320+ # Gambaterm configuration
321+ gambaterm_config_dir = Path (
322+ os .environ .get ("GAMBATERM_CONFIG_DIR" , "~/.config/gambaterm" )
323+ ).expanduser ()
324+ server_host_key = gambaterm_config_dir / "ssh_host_key"
325+ config_authorized_keys = gambaterm_config_dir / "authorized_keys"
326+
327+ # User SSH public keys (for authentication)
328+ user_ssh_dir = Path (os .environ .get ("GAMBATERM_USER_SSH_DIR" , "~/.ssh" )).expanduser ()
329+ user_authorized_keys = user_ssh_dir / "authorized_keys"
330+
331+ # Generate host key if it does not exist
332+ if not server_host_key .exists ():
333+ print (f"Generating SSH host key at { server_host_key } ..." )
334+ server_host_key .parent .mkdir (parents = True , exist_ok = True )
335+ key = asyncssh .generate_private_key ("ssh-ed25519" )
336+ server_host_key .write_bytes (key .export_private_key ())
337+ server_host_key .chmod (0o600 )
338+ server_host_keys = [str (server_host_key )]
339+
340+ # Collect authorized client keys for public key authentication
341+ authorized_client_keys = []
342+ if isinstance (
343+ authentication , (PublicKeyAuthentication , PasswordAndPublicKeyAuthentication )
344+ ):
345+ for key_type in ["rsa" , "ed25519" , "ecdsa" ]:
346+ user_public_key = user_ssh_dir / f"id_{ key_type } .pub"
347+ if user_public_key .exists ():
348+ authorized_client_keys .append (str (user_public_key ))
349+ if user_authorized_keys .exists ():
350+ authorized_client_keys .append (str (user_authorized_keys ))
351+ if config_authorized_keys .exists ():
352+ authorized_client_keys .append (str (config_authorized_keys ))
353+ if not authorized_client_keys and isinstance (
354+ authentication , PublicKeyAuthentication
355+ ):
300356 raise SystemExit (
301- f"The server requires a private RSA key to use as a host hey.\n "
302- f"You may generate one by running the following command:\n \n "
303- f" ssh-keygen -f { ssh_key_dir / 'id_rsa' } -P ''\n "
357+ f"Public key authentication is enabled, but no authorized keys were found.\n "
358+ f"Please add the public keys of allowed clients to { config_authorized_keys } ."
304359 )
305- server_host_keys = [str (user_private_key )]
306- authorized_client_keys = []
307- if user_public_key .exists ():
308- authorized_client_keys = [str (user_public_key )]
309360
310361 # Remove chacha20 from encryption_algs because it's a bit too expensive
311362 encryption_algs = [
@@ -318,7 +369,7 @@ async def run_server(
318369 ]
319370
320371 server = await asyncssh .create_server (
321- lambda : SSHServer (password , console_cls , namespace , executor ),
372+ lambda : SSHServer (authentication , console_cls , namespace , executor ),
322373 bind ,
323374 port ,
324375 server_host_keys = server_host_keys ,
@@ -328,8 +379,22 @@ async def run_server(
328379 line_editor = False ,
329380 reuse_address = True ,
330381 )
382+
383+ match authentication :
384+ case NoAuthentication ():
385+ print ("Authentication disabled (no password nor public key required)" )
386+ case PasswordAndPublicKeyAuthentication ():
387+ print ("Authentication methods:" )
388+ print ("- Global password" )
389+ for key_path in authorized_client_keys :
390+ print (f"- Public keys from: { key_path } " )
391+ case PublicKeyAuthentication ():
392+ print ("Authentication methods:" )
393+ for key_path in authorized_client_keys :
394+ print (f"- Public keys from: { key_path } " )
331395 bind , port = server .sockets [0 ].getsockname ()
332- print (f"Running ssh server on { bind } :{ port } ..." , flush = True )
396+ print (f"Running SSH server on { bind } :{ port } ..." , flush = True )
397+
333398 async with server :
334399 # Sleep forever
335400 await asyncio .Future ()
@@ -365,19 +430,37 @@ def main(
365430 default = None ,
366431 help = "Enable password authentification with the given global password" ,
367432 )
433+ parser .add_argument (
434+ "--no-auth" ,
435+ action = "store_true" ,
436+ help = "Disable authentication altogether (no password nor public key required)" ,
437+ )
368438
369439 # Parse arguments
370440 namespace = parser .parse_args (parser_args )
371441 bind : str = namespace .__dict__ .pop ("bind" )
372442 port : int = namespace .__dict__ .pop ("port" )
373443 password : str = namespace .__dict__ .pop ("password" )
444+ no_auth : bool = namespace .__dict__ .pop ("no_auth" )
445+
446+ # Determine authentication method
447+ if no_auth and password is None :
448+ authentication : AuthenticationMethod = NoAuthentication ()
449+ elif not no_auth and password is not None :
450+ authentication = PasswordAndPublicKeyAuthentication (password )
451+ elif not no_auth and password is None :
452+ authentication = PublicKeyAuthentication ()
453+ else :
454+ raise SystemExit (
455+ "Both `--password` and `--no-auth` cannot be provided at the same time"
456+ )
374457
375458 # Run an executor with no limit on the number of threads
376459 try :
377460 with ThreadPoolExecutor (max_workers = 32 ) as executor :
378461 # Run the server in asyncio
379462 asyncio .run (
380- run_server (bind , port , password , console_cls , namespace , executor )
463+ run_server (bind , port , authentication , console_cls , namespace , executor )
381464 )
382465 except KeyboardInterrupt :
383466 pass
0 commit comments