| ''' |
| Copyright 2019 Google LLC |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| https://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ''' |
| |
| |
| import re |
| import os |
| import sys |
| |
| from mdt import command |
| from mdt import console |
| from mdt import keys |
| |
| |
| class ShellCommand(command.NetworkCommand): |
| '''Usage: mdt shell [<device-or-ip-address>] |
| |
| Opens an interactive shell to either your preferred device or to the first |
| device found. |
| |
| If <device-or-ip-address> is specified, shell will attempt to connect to that |
| device name or directly to the IP address provided instead. |
| |
| Variables used: |
| preferred-device - set this to your preferred device name to connect |
| to by default if no <devicename> is provided on the |
| command line. Can be set to an IPv4 address to bypass |
| the mDNS lookup. |
| username - set this to the username that should be used to |
| connect to a device with. Defaults to 'mendel'. |
| password - set this to the password to use to login to a new |
| device with. Defaults to 'mendel'. Only used |
| during the initial setup phase of pushing an SSH |
| key to the board. |
| |
| If no SSH key is available on disk (ie: you didn't run genkey before running |
| shell), this will implicitly run genkey for you. Additionally, shell will |
| attempt to connect to a device by doing the following: |
| |
| 1. Attempt a connection using your SSH key only, with no password. |
| 2. If the connection attempt failed due to authentication, will |
| attempt to push the key to the device by using the default |
| login credentials in the 'username' and 'password' variables. |
| 3. Installs your SSH key to the device after logging in. |
| 4. Disconnects and reconnects using the SSH key. |
| ''' |
| |
| def preConnectRun(self, args): |
| if len(args) > 2: |
| print("Usage: mdt shell [<device-or-ip-address>]") |
| return False |
| |
| if len(args) == 2: |
| self.device = args[1] |
| |
| return True |
| |
| def runWithClient(self, client, args): |
| channel = client.openShell() |
| cons = console.Console(channel, sys.stdin) |
| return cons.run() |
| |
| |
| class ExecCommand(command.NetworkCommand): |
| '''Usage: mdt exec [<shell-command...>] |
| |
| Opens a non-interactive shell to either your preferred device or to the first |
| device found. |
| |
| Variables used: |
| preferred-device - set this to your preferred device name to connect |
| to by default if no <devicename> is provided on the |
| command line. Can be set to an IPv4 address to bypass |
| the mDNS lookup. |
| username - set this to the username that should be used to |
| connect to a device with. Defaults to 'mendel'. |
| password - set this to the password to use to login to a new |
| device with. Defaults to 'mendel'. Only used |
| during the initial setup phase of pushing an SSH |
| key to the board. |
| |
| If no SSH key is available on disk (ie: you didn't run genkey before running |
| shell), this will implicitly run genkey for you. Additionally, shell will |
| attempt to connect to a device by doing the following: |
| |
| 1. Attempt a connection using your SSH key only, with no password. |
| 2. If the connection attempt failed due to authentication, will |
| attempt to push the key to the device by using the default |
| login credentials in the 'username' and 'password' variables. |
| 3. Installs your SSH key to the device after logging in. |
| 4. Disconnects and reconnects using the SSH key. |
| ''' |
| def runWithClient(self, client, args): |
| channel = client.shellExec(' '.join(args[1:]), allocPty=True) |
| cons = console.Console(channel, sys.stdin) |
| return cons.run() |
| |
| |
| class RebootCommand(command.NetworkCommand): |
| def runWithClient(self, client, args): |
| channel = client.shellExec("sudo reboot", allocPty=True) |
| cons = console.Console(channel, sys.stdin) |
| return cons.run() |
| |
| |
| class RebootBootloaderCommand(command.NetworkCommand): |
| def runWithClient(self, client, args): |
| channel = client.shellExec("sudo reboot-bootloader", allocPty=True) |
| cons = console.Console(channel, sys.stdin) |
| return cons.run() |
| |
| |
| class PushKeyCommand(command.NetworkCommand): |
| '''Usage: mdt pushkey [<path-to-ssh-public-key>] |
| |
| Copies an SSH public key provided to the device's ~/.ssh/authorized_keys |
| file. If no public key is provided, attempts to push MDTs previously generated |
| public key from ~/.config/mdt/keys/mdt.key. |
| ''' |
| |
| def _pushMdtKey(self, client): |
| print('Pushing {0}'.format(keys.KEYFILE_PATH)) |
| client.pushKey() |
| print('Push complete.') |
| return 0 |
| |
| def _pushOtherKey(self, client, keyfile): |
| sftp = client.openSftp() |
| |
| if not os.path.exists(source_keyfile): |
| print("Can't copy {0}: no such file or directory.".format(source_keyfile)) |
| return 1 |
| |
| source_key = '' |
| with open(args[1], 'rb') as fp: |
| source_key = fp.read() |
| |
| try: |
| sftp.chdir('/home/mendel/.ssh') |
| except FileNotFoundError as e: |
| sftp.mkdir('/home/mendel/.ssh', mode=0o700) |
| |
| with sftp.open('/home/mendel/.ssh/authorized_keys', 'a+b') as fp: |
| fp.write('\r\n') |
| fp.write(source_key) |
| |
| print("Key {0} pushed.".format(source_keyfile)) |
| return 0 |
| |
| def runWithClient(self, client, args): |
| key_to_push = None |
| |
| # No arguments? Let the usual client push take effect. |
| if len(args) == 1: |
| return self._pushMdtKey(client) |
| |
| source_keyfile = args[1] |
| print('Pushing {0}'.format(source_keyfile)) |
| return self._pushOtherKey(client, source_keyfile) |
| |
| |
| class ResetKeysCommand(command.NetworkCommand): |
| '''Usage: mdt resetkeys <device-or-ip-address> |
| |
| Resets a device to it's pre-MDT state by removing all MDT keys and restarting |
| the mdt-keymaster on the device so that new keys can be pushed again.''' |
| |
| def preConnectRun(self, args): |
| if len(args) != 2: |
| print("Usage: mdt resetkeys <device-or-ip-address>") |
| return False |
| |
| if len(args) == 2: |
| self.device = args[1] |
| |
| return True |
| |
| def runWithClient(self, client, args): |
| # Setup this session now, since once we remove the key from |
| # authorized_keys, we won't be able to use execSession. |
| channel = client.openChannel() |
| |
| sftp = client.openSftp() |
| try: |
| sftp.chdir('/home/mendel/.ssh') |
| except FileNotFoundError as e: |
| print('No keys were previously pushed to the board.') |
| else: |
| lines = [] |
| |
| with sftp.open('/home/mendel/.ssh/authorized_keys', 'r') as fp: |
| lines = fp.readlines() |
| |
| with sftp.open('/home/mendel/.ssh/authorized_keys', 'w') as fp: |
| for line in lines: |
| if ' mdt' not in line: |
| print('wrote: {0}'.format(line)) |
| fp.write(line) |
| |
| channel.exec_command("sudo systemctl restart mdt-keymaster") |
| cons = console.Console(channel, sys.stdin) |
| try: |
| cons.run() |
| except console.ConnectionClosedError as e: |
| if e.exit_code: |
| print('`systemctl restart mdt-keymaster` exited with code {0}'.format(e.exit_code)) |
| print('Your device may be in an inconsistent state. Verify using') |
| print('the serial console.') |
| else: |
| print('Successfully reset {0}'.format(self.device)) |
| return e.exit_code |