CAR/wly_experiment/dta_codes/manager/Machine.py
2025-04-14 22:25:58 +08:00

81 lines
1.9 KiB
Python

#!/usr/bin/env python3
#This script contains the Machine class, with various functions shared between the components
import time
import pexpect
from common import log
class Machine:
host = None
name = None
def logger(self, text):
log("%s: \t%s" %(self.name, text))
#High verbosity output
def debug(self, text):
self.logger(" Debug: %s" %(text))
def error(self, text):
self.logger("ERROR: %s" %(text))
def reboot(self):
self.logger("Rebooting %s at %s" %(self.name, self.host))
ssh = self.init_ssh()
ssh.sendline("sudo reboot")
i = ssh.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=10)
assert i == 0, "Failed to detect reboot!"
time.sleep(1)
def testConnection(self):
self.debug("Testing connection to %s at %s..." %(self.name, self.host))
p = pexpect.spawn("ssh %s" %self.host)
i = p.expect(["Welcome to Ubuntu", "Connection refused", pexpect.TIMEOUT, pexpect.EOF], timeout=5)
if i == 0:
self.debug("Logged into %s" %self.host)
elif i == 1:
self.debug("Connection refused!")
return False
elif i == 2:
self.debug("Connection timeout!")
return False
elif i == 3:
self.debug("SSH terminates!")
return False
self.debug("Verifying command capability...")
content = "Testing"
p.sendline("echo \"%s\" > ssh_works" %content)
p.expect("$")
p.sendline("cat ssh_works")
i = p.expect([content, pexpect.TIMEOUT], timeout=2)
if i != 0:
self.error("Did not find expected output!")
return False
self.debug("Commands work! Resetting and logging out...")
p.sendline("rm ssh_works")
p.expect("$")
p.sendline("exit")
time.sleep(1)
p.close()
return True
def init_ssh(self):
self.debug("Logging into %s at %s..." %(self.name, self.host))
p = pexpect.spawn("ssh %s" %self.host)
i = p.expect(["$", pexpect.TIMEOUT], timeout=5)
if i != 0:
self.error("Timeout!")
return None
self.debug("SSH to %s is initiated" %self.host)
return p