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 | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 25 | self.client.connect(self.host, username = self.user, password = self.password, |
A R Karthick | 9313b76 | 2016-11-07 13:14:35 -0800 | [diff] [blame] | 26 | key_filename = self.key_file, timeout=timeout, port = self.port) |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 27 | except: |
| 28 | log.error('Unable to connect to test host %s' %self.host) |
| 29 | return False, None |
| 30 | |
| 31 | channel = self.client.get_transport().open_session() |
| 32 | channel.exec_command(cmd) |
A R Karthick | 9313b76 | 2016-11-07 13:14:35 -0800 | [diff] [blame] | 33 | if channel.exit_status_ready(): |
| 34 | status = channel.recv_exit_status() |
| 35 | else: |
| 36 | status = 0 |
A R Karthick | ec1dde0 | 2016-10-06 13:52:25 -0700 | [diff] [blame] | 37 | output = None |
| 38 | st = status == 0 |
| 39 | if st: |
| 40 | output = '' |
| 41 | while True: |
| 42 | data = channel.recv(4096) |
| 43 | if data: |
| 44 | output += data |
| 45 | else: |
| 46 | break |
| 47 | channel.close() |
| 48 | return st, output |
| 49 | |
| 50 | if __name__ == '__main__': |
| 51 | agent = SSHTestAgent(user = 'ubuntu', password = 'ubuntu') |
| 52 | cmds = ('docker images', 'docker ps') |
| 53 | for cmd in cmds: |
| 54 | st, output = agent.run_cmd(cmd) |
| 55 | print('Command \"%s\" returned with status: %s' %(cmd, st)) |
| 56 | if st: |
| 57 | print('%s\n' %output) |