A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 1 | import os, sys |
| 2 | from paramiko import SSHClient, WarningPolicy, AutoAddPolicy |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 3 | from scapy.all import * |
| 4 | |
| 5 | class SSHTestAgent(object): |
| 6 | key_file = os.getenv('SSH_KEY_FILE', None) |
A R Karthick | 19aaf5c | 2016-11-09 17:47:57 -0800 | [diff] [blame] | 7 | host = os.getenv('CORD_TEST_HOST', '172.17.0.1') |
| 8 | hosts_file = os.path.join(os.getenv('HOME'), '.ssh', 'known_hosts') |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 9 | user = 'ubuntu' |
| 10 | password = None |
| 11 | |
A R Karthick | 9313b76 | 2016-11-07 13:14:35 -0800 | [diff] [blame] | 12 | def __init__(self, host = host, user = user, password = password, port = 22): |
| 13 | self.host = host |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 14 | self.user = user |
| 15 | self.password = password |
A R Karthick | 9313b76 | 2016-11-07 13:14:35 -0800 | [diff] [blame] | 16 | self.port = port |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 17 | self.client = SSHClient() |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 18 | self.client.set_missing_host_key_policy(AutoAddPolicy()) |
| 19 | |
| 20 | def run_cmd(self, cmd, timeout = 5): |
| 21 | """Run the command on the test host""" |
A R Karthick | 19aaf5c | 2016-11-09 17:47:57 -0800 | [diff] [blame] | 22 | host_remove = 'ssh-keygen -f "%s" -R [%s]:8101' %(self.hosts_file, self.host) |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 23 | try: |
A R Karthick | 19aaf5c | 2016-11-09 17:47:57 -0800 | [diff] [blame] | 24 | os.system(host_remove) |
A R Karthick | d0a334d | 2016-11-10 17:47:08 -0800 | [diff] [blame] | 25 | except: pass |
| 26 | |
| 27 | try: |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 28 | self.client.connect(self.host, username = self.user, password = self.password, |
A R Karthick | 9313b76 | 2016-11-07 13:14:35 -0800 | [diff] [blame] | 29 | key_filename = self.key_file, timeout=timeout, port = self.port) |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 30 | except: |
| 31 | log.error('Unable to connect to test host %s' %self.host) |
| 32 | return False, None |
| 33 | |
| 34 | channel = self.client.get_transport().open_session() |
| 35 | channel.exec_command(cmd) |
A R Karthick | 9313b76 | 2016-11-07 13:14:35 -0800 | [diff] [blame] | 36 | if channel.exit_status_ready(): |
| 37 | status = channel.recv_exit_status() |
| 38 | else: |
| 39 | status = 0 |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 40 | output = None |
| 41 | st = status == 0 |
| 42 | if st: |
| 43 | output = '' |
| 44 | while True: |
| 45 | data = channel.recv(4096) |
| 46 | if data: |
| 47 | output += data |
| 48 | else: |
| 49 | break |
| 50 | channel.close() |
A R Karthick | d0a334d | 2016-11-10 17:47:08 -0800 | [diff] [blame] | 51 | self.client.close() |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 52 | return st, output |
| 53 | |
| 54 | if __name__ == '__main__': |
| 55 | agent = SSHTestAgent(user = 'ubuntu', password = 'ubuntu') |
| 56 | cmds = ('docker images', 'docker ps') |
| 57 | for cmd in cmds: |
| 58 | st, output = agent.run_cmd(cmd) |
| 59 | print('Command \"%s\" returned with status: %s' %(cmd, st)) |
| 60 | if st: |
| 61 | print('%s\n' %output) |