1
1
# pylint: disable=unused-argument
2
2
"""The ShellDriver provides the CommandProtocol, ConsoleProtocol and
3
3
InfoProtocol on top of a SerialPort."""
4
+ import os
4
5
import io
5
6
import re
6
7
import shlex
@@ -34,6 +35,7 @@ class ShellDriver(CommandMixin, Driver, CommandProtocol, FileTransferProtocol):
34
35
username (str): username to login with
35
36
password (str): password to login with
36
37
keyfile (str): keyfile to bind mount over users authorized keys
38
+ dest_authorized_keys (str): optional, default="~/.ssh/authorized_keys", filename of the authorized_keys file
37
39
login_timeout (int): optional, timeout for login prompt detection
38
40
console_ready (regex): optional, pattern used by the kernel to inform the user that a
39
41
console can be activated by pressing enter.
@@ -49,6 +51,7 @@ class ShellDriver(CommandMixin, Driver, CommandProtocol, FileTransferProtocol):
49
51
username = attr .ib (validator = attr .validators .instance_of (str ))
50
52
password = attr .ib (default = None , validator = attr .validators .optional (attr .validators .instance_of (str )))
51
53
keyfile = attr .ib (default = "" , validator = attr .validators .instance_of (str ))
54
+ dest_authorized_keys = attr .ib (default = "~/.ssh/authorized_keys" , validator = attr .validators .instance_of (str ))
52
55
login_timeout = attr .ib (default = 60 , validator = attr .validators .instance_of (int ))
53
56
console_ready = attr .ib (default = "" , validator = attr .validators .instance_of (str ))
54
57
await_login_timeout = attr .ib (default = 2 , validator = attr .validators .instance_of (int ))
@@ -72,7 +75,7 @@ def on_activate(self):
72
75
if self .target .env :
73
76
keyfile_path = self .target .env .config .resolve_path (self .keyfile )
74
77
75
- self ._put_ssh_key (keyfile_path )
78
+ self ._put_ssh_key (keyfile_path , self . dest_authorized_keys )
76
79
77
80
def on_deactivate (self ):
78
81
self ._status = 0
@@ -217,8 +220,8 @@ def _write_key(self, keyline, dest):
217
220
self ._run_check (f'echo -n "{ part } " >> { dest } ' )
218
221
self ._run_check (f'echo "" >> { dest } ' )
219
222
220
- @step (args = ['keyfile_path' ])
221
- def _put_ssh_key (self , keyfile_path ):
223
+ @step (args = ['keyfile_path' , 'dest_authorized_keys' ])
224
+ def _put_ssh_key (self , keyfile_path , dest_authorized_keys ):
222
225
"""Upload an SSH Key to a target"""
223
226
regex = re .compile (
224
227
r"""ssh-(rsa|ed25519)
@@ -236,7 +239,8 @@ def _put_ssh_key(self, keyfile_path):
236
239
f"Could not parse SSH-Key from file: { keyfile } "
237
240
)
238
241
self .logger .debug ("Read Key: %s" , new_key )
239
- auth_keys , _ , read_keys = self ._run ("cat ~/.ssh/authorized_keys" )
242
+ dest_authorized_keys_dir = os .path .dirname (dest_authorized_keys )
243
+ auth_keys , _ , read_keys = self ._run (f"""cat { self .dest_authorized_keys } """ )
240
244
self .logger .debug ("Exitcode trying to read keys: %s, keys: %s" , read_keys , auth_keys )
241
245
result = []
242
246
_ , _ , test_write = self ._run ("touch ~/.test" )
@@ -258,35 +262,37 @@ def _put_ssh_key(self, keyfile_path):
258
262
259
263
if test_write == 0 and read_keys == 0 :
260
264
self .logger .debug ("Key not on target and writeable, concatenating..." )
261
- self ._write_key (keyline , "~/.ssh/authorized_keys" )
265
+ self ._write_key (keyline , dest_authorized_keys )
262
266
self ._run_check ("rm ~/.test" )
263
267
return
264
268
265
269
if test_write == 0 :
266
- self .logger .debug ("Key not on target, testing for .ssh directory" )
267
- _ , _ , ssh_dir = self ._run (" [ -d ~/.ssh/ ] " )
270
+ self .logger .debug ("Key not on target, testing for % directory" , dest_authorized_keys_dir )
271
+ _ , _ , ssh_dir = self ._run (f""" [ -d { dest_authorized_keys_dir } ]"" " )
268
272
if ssh_dir != 0 :
269
- self .logger .debug ("~/.ssh did not exist, creating" )
270
- self ._run ("mkdir ~/.ssh/" )
271
- self ._run_check ("chmod 700 ~/.ssh/" )
272
- self .logger .debug ("Creating ~/.ssh/authorized_keys" )
273
- self ._run_check ("touch ~/.ssh/authorized_keys" )
274
- self ._write_key (keyline , "~/.ssh/authorized_keys" )
273
+ self .logger .debug (" % did not exits, creating" , dest_authorized_keys_dir )
274
+ self ._run (f"""mkdir -p { dest_authorized_keys_dir } """ )
275
+ self ._run_check (f"""chmod 700 { dest_authorized_keys_dir } """ )
276
+ self .logger .debug ("Creating %" , dest_authorized_keys )
277
+ self ._write_key (keyline , dest_authorized_keys )
275
278
self ._run_check ("rm ~/.test" )
276
279
return
277
280
278
281
self .logger .debug ("Key not on target and not writeable, using bind mount..." )
279
282
self ._run_check ('mkdir -p -m 700 /tmp/labgrid-ssh/' )
280
- self ._run (" cp -a ~/.ssh/ * /tmp/labgrid-ssh/" )
283
+ self ._run (f""" cp -a { dest_authorized_keys_dir } / * /tmp/labgrid-ssh/"" " )
281
284
self ._write_key (keyline , "/tmp/labgrid-ssh/authorized_keys" )
282
285
self ._run_check ('chmod 600 /tmp/labgrid-ssh/authorized_keys' )
283
- out , err , exitcode = self ._run (' mount --bind /tmp/labgrid-ssh/ ~/.ssh/' )
286
+ out , err , exitcode = self ._run (f""" mount --bind /tmp/labgrid-ssh/ { dest_authorized_keys_dir } """ )
284
287
if exitcode != 0 :
285
- self .logger .warning ("Could not bind mount ~/.ssh directory: %s %s" , out , err )
288
+ self .logger .warning ("Could not bind mount % directory: % %" ,
289
+ dest_authorized_keys_dir , out , err )
286
290
287
291
@Driver .check_active
288
- def put_ssh_key (self , keyfile_path ):
289
- self ._put_ssh_key (keyfile_path )
292
+ def put_ssh_key (self , keyfile_path , dest_authorized_keys = None ):
293
+ if dest_authorized_keys is None :
294
+ dest_authorized_keys = self .dest_authorized_keys
295
+ self ._put_ssh_key (keyfile_path , dest_authorized_keys )
290
296
291
297
def _xmodem_getc (self , size , timeout = 10 ):
292
298
""" called by the xmodem.XMODEM instance to read protocol data from the console """
0 commit comments