gem5: Addition of a gem5 connection

This is the first in a series of commits to allow for the use of gem5 as a target in devlib.
The interaction with the gem5 simulation is done by the Gem5Connection, a type of TelnetConnection.
There are two types of Gem5Connections (Linux and Android) as they differ slightly in the implementation
of some methods. Other types of connections have only been changed to make them uniform (e.g. arguments for init method)

 Changes to be committed:
	modified:   devlib/__init__.py
	modified:   devlib/host.py
	modified:   devlib/utils/android.py
	modified:   devlib/utils/ssh.py

4 files changed, 519 insertions, 13 deletions

diff --git a/devlib/__init__.py b/devlib/__init__.py
index 86949d3..f4ff0f3 100644
--- a/devlib/__init__.py
+++ b/devlib/__init__.py
@@ -19,4 +19,4 @@ from devlib.trace.ftrace import FtraceCollector
 
 from devlib.host import LocalConnection
 from devlib.utils.android import AdbConnection
-from devlib.utils.ssh import SshConnection, TelnetConnection
+from devlib.utils.ssh import SshConnection, TelnetConnection, Gem5Connection
diff --git a/devlib/host.py b/devlib/host.py
index 121e5b4..3e42c0f 100644
--- a/devlib/host.py
+++ b/devlib/host.py
@@ -49,7 +49,8 @@ class LocalConnection(object):
         else:
             shutil.copy(source, dest)
 
-    def execute(self, command, timeout=None, check_exit_code=True, as_root=False):
+    def execute(self, command, timeout=None, check_exit_code=True,
+                as_root=False, strip_colors=True):
         self.logger.debug(command)
         if as_root:
             if self.unrooted:
diff --git a/devlib/utils/android.py b/devlib/utils/android.py
index b3479a8..ad4c324 100644
--- a/devlib/utils/android.py
+++ b/devlib/utils/android.py
@@ -186,7 +186,7 @@ class AdbConnection(object):
             self.ls_command = 'ls'
         logger.info("ls command is set to {}".format(self.ls_command))
 
-    def __init__(self, device=None, timeout=None):
+    def __init__(self, device=None, timeout=None, platform=None):
         self.timeout = timeout if timeout is not None else self.default_timeout
         if device is None:
             device = adb_get_device(timeout=timeout)
@@ -216,7 +216,8 @@ class AdbConnection(object):
         command = "pull '{}' '{}'".format(source, dest)
         return adb_command(self.device, command, timeout=timeout)
 
-    def execute(self, command, timeout=None, check_exit_code=False, as_root=False):
+    def execute(self, command, timeout=None, check_exit_code=False,
+                as_root=False, strip_colors=True):
         return adb_shell(self.device, command, timeout, check_exit_code,
                          as_root, self.newline_separator)
 
diff --git a/devlib/utils/ssh.py b/devlib/utils/ssh.py
index 448f9e7..ff30346 100644
--- a/devlib/utils/ssh.py
+++ b/devlib/utils/ssh.py
@@ -22,6 +22,7 @@ import re
 import threading
 import tempfile
 import shutil
+import socket
 import time
 
 import pexpect
@@ -34,14 +35,16 @@ from pexpect import EOF, TIMEOUT, spawn
 
 from devlib.exception import HostError, TargetError, TimeoutError
 from devlib.utils.misc import which, strip_bash_colors, escape_single_quotes, check_output
+from devlib.utils.types import boolean
 
 
 ssh = None
 scp = None
 sshpass = None
 
-logger = logging.getLogger('ssh')
 
+logger = logging.getLogger('ssh')
+gem5_logger = logging.getLogger('gem5-connection')
 
 def ssh_get_shell(host, username, password=None, keyfile=None, port=None, timeout=10, telnet=False, original_prompt=None):
     _check_env()
@@ -91,18 +94,20 @@ class TelnetPxssh(pxssh.pxssh):
 
         spawn._spawn(self, cmd)  # pylint: disable=protected-access
 
-        if password is None:
-            i = self.expect([self.original_prompt, 'Login timed out'], timeout=login_timeout)
-        else:
+        try:
             i = self.expect('(?i)(?:password)', timeout=login_timeout)
             if i == 0:
                 self.sendline(password)
                 i = self.expect([self.original_prompt, 'Login incorrect'], timeout=login_timeout)
-            else:
-                raise pxssh.ExceptionPxssh('could not log in: did not see a password prompt')
-
             if i:
                 raise pxssh.ExceptionPxssh('could not log in: password was incorrect')
+        except TIMEOUT:
+            if not password:
+                # No password promt before TIMEOUT & no password provided
+                # so assume everything is okay
+                pass
+            else:
+                raise pxssh.ExceptionPxssh('could not log in: did not see a password prompt')
 
         if not self.sync_original_prompt(sync_multiplier):
             self.close()
@@ -155,6 +160,7 @@ class SshConnection(object):
                  telnet=False,
                  password_prompt=None,
                  original_prompt=None,
+                 platform=None
                  ):
         self.host = host
         self.username = username
@@ -175,7 +181,8 @@ class SshConnection(object):
         source = '{}@{}:{}'.format(self.username, self.host, source)
         return self._scp(source, dest, timeout)
 
-    def execute(self, command, timeout=None, check_exit_code=True, as_root=False, strip_colors=True):
+    def execute(self, command, timeout=None, check_exit_code=True,
+                as_root=False, strip_colors=True): #pylint: disable=unused-argument
         try:
             with self.lock:
                 output = self._execute_and_wait_for_prompt(command, timeout, as_root, strip_colors)
@@ -286,7 +293,7 @@ class TelnetConnection(SshConnection):
                  timeout=None,
                  password_prompt=None,
                  original_prompt=None,
-                 ):
+                 platform=None):
         self.host = host
         self.username = username
         self.password = password
@@ -299,6 +306,503 @@ class TelnetConnection(SshConnection):
         self.conn = ssh_get_shell(host, username, password, None, port, timeout, True, original_prompt)
 
 
+class Gem5Connection(TelnetConnection):
+
+    def __init__(self,
+                 platform,
+                 host=None,
+                 username=None,
+                 password=None,
+                 port=None,
+                 timeout=None,
+                 password_prompt=None,
+                 original_prompt=None,
+                 ):
+        if host is not None:
+            host_system = socket.gethostname()
+            if host_system != host:
+                raise TargetError("Gem5Connection can only connect to gem5 "
+                                   "simulations on your current host, which "
+                                   "differs from the one given {}!"
+                                   .format(host_system, host))
+        if username is not None and username != 'root':
+            raise ValueError('User should be root in gem5!')
+        if password is not None and password != '':
+            raise ValueError('No password needed in gem5!')
+        self.username = 'root'
+        self.is_rooted = True
+        self.password = None
+        self.port = None
+        # Long timeouts to account for gem5 being slow
+        # Can be overriden if the given timeout is longer
+        self.default_timeout = 3600
+        if timeout is not None:
+            if timeout > self.default_timeout:
+                logger.info('Overwriting the default timeout of gem5 ({})'
+                                 ' to {}'.format(self.default_timeout, timeout))
+                self.default_timeout = timeout
+            else:
+                logger.info('Ignoring the given timeout --> gem5 needs longer timeouts')
+        self.ready_timeout = self.default_timeout * 3
+        # Counterpart in gem5_interact_dir
+        self.gem5_input_dir = '/mnt/host/'
+        # Location of m5 binary in the gem5 simulated system
+        self.m5_path = None
+        # Actual telnet connection to gem5 simulation
+        self.conn = None
+        # Flag to indicate the gem5 device is ready to interact with the
+        # outer world
+        self.ready = False
+        # Lock file to prevent multiple connections to same gem5 simulation
+        # (gem5 does not allow this)
+        self.lock_directory = '/tmp/'
+        self.lock_file_name = None # Will be set once connected to gem5
+
+        # These parameters will be set by either the method to connect to the
+        # gem5 platform or directly to the gem5 simulation
+        # Intermediate directory to push things to gem5 using VirtIO
+        self.gem5_interact_dir = None
+        # Directory to store output  from gem5 on the host
+        self.gem5_out_dir = None
+        # Actual gem5 simulation
+        self.gem5simulation = None
+
+        # Connect to gem5
+        if platform:
+            self._connect_gem5_platform(platform)
+
+        # Wait for boot
+        self._wait_for_boot()
+
+        # Mount the virtIO to transfer files in/out gem5 system
+        self._mount_virtio()
+
+    def set_hostinteractdir(self, indir):
+        logger.info('Setting hostinteractdir  from {} to {}'
+                    .format(self.gem5_input_dir, indir))
+        self.gem5_input_dir = indir
+
+    def push(self, source, dest, timeout=None):
+        """
+        Push a file to the gem5 device using VirtIO
+
+        The file to push to the device is copied to the temporary directory on
+        the host, before being copied within the simulation to the destination.
+        Checks, in the form of 'ls' with error code checking, are performed to
+        ensure that the file is copied to the destination.
+        """
+        # First check if the connection is set up to interact with gem5
+        self._check_ready()
+
+        filename = os.path.basename(source)
+        logger.debug("Pushing {} to device.".format(source))
+        logger.debug("gem5interactdir: {}".format(self.gem5_interact_dir))
+        logger.debug("dest: {}".format(dest))
+        logger.debug("filename: {}".format(filename))
+
+        # We need to copy the file to copy to the temporary directory
+        self._move_to_temp_dir(source)
+
+        # Dest in gem5 world is a file rather than directory
+        if os.path.basename(dest) != filename:
+            dest = os.path.join(dest, filename)
+        # Back to the gem5 world
+        self._gem5_shell("ls -al {}{}".format(self.gem5_input_dir, filename))
+        self._gem5_shell("cat '{}''{}' > '{}'".format(self.gem5_input_dir,
+                                                     filename,
+                                                     dest))
+        self._gem5_shell("sync")
+        self._gem5_shell("ls -al {}".format(dest))
+        self._gem5_shell("ls -al {}".format(self.gem5_input_dir))
+        logger.debug("Push complete.")
+
+    def pull(self, source, dest, timeout=0): #pylint: disable=unused-argument
+        """
+        Pull a file from the gem5 device using m5 writefile
+
+        The file is copied to the local directory within the guest as the m5
+        writefile command assumes that the file is local. The file is then
+        written out to the host system using writefile, prior to being moved to
+        the destination on the host.
+        """
+        # First check if the connection is set up to interact with gem5
+        self._check_ready()
+
+        filename = os.path.basename(source)
+
+        logger.debug("pull_file {} {}".format(source, filename))
+        # We don't check the exit code here because it is non-zero if the source
+        # and destination are the same. The ls below will cause an error if the
+        # file was not where we expected it to be.
+        if os.path.dirname(source) != os.getcwd():
+            self._gem5_shell("cat '{}' > '{}'".format(source, filename))
+        self._gem5_shell("sync")
+        self._gem5_shell("ls -la {}".format(filename))
+        logger.debug('Finished the copy in the simulator')
+        self._gem5_util("writefile {}".format(filename))
+
+        if 'cpu' not in filename:
+            while not os.path.exists(os.path.join(self.gem5_out_dir, filename)):
+                time.sleep(1)
+
+        # Perform the local move
+        shutil.move(os.path.join(self.gem5_out_dir, filename), dest)
+        logger.debug("Pull complete.")
+
+    def execute(self, command, timeout=1000, check_exit_code=True,
+                as_root=False, strip_colors=True):
+        """
+        Execute a command on the gem5 platform
+        """
+        # First check if the connection is set up to interact with gem5
+        self._check_ready()
+
+        output = self._gem5_shell(command, as_root=as_root)
+        if strip_colors:
+            output = strip_bash_colors(output)
+        return output
+
+    def background(self, command, stdout=subprocess.PIPE,
+                   stderr=subprocess.PIPE, as_root=False):
+        # First check if the connection is set up to interact with gem5
+        self._check_ready()
+
+        # Create the logfile for stderr/stdout redirection
+        command_name = command.split(' ')[0].split('/')[-1]
+        redirection_file = 'BACKGROUND_{}.log'.format(command_name)
+        trial = 0
+        while os.path.isfile(redirection_file):
+            # Log file already exists so add to name
+           redirection_file = 'BACKGROUND_{}{}.log'.format(command_name, trial)
+           trial += 1
+
+        # Create the command to pass on to gem5 shell
+        complete_command = '{} >> {} 2>&1 &'.format(command, redirection_file)
+        output = self._gem5_shell(complete_command, as_root=as_root)
+        output = strip_bash_colors(output)
+        gem5_logger.info('STDERR/STDOUT of background command will be '
+                         'redirected to {}. Use target.pull() to '
+                         'get this file'.format(redirection_file))
+        return output
+
+    def close(self):
+        """
+        Close and disconnect from the gem5 simulation. Additionally, we remove
+        the temporary directory used to pass files into the simulation.
+        """
+        gem5_logger.info("Gracefully terminating the gem5 simulation.")
+        try:
+            self._gem5_util("exit")
+            self.gem5simulation.wait()
+        except EOF:
+            pass
+        gem5_logger.info("Removing the temporary directory")
+        try:
+            shutil.rmtree(self.gem5_interact_dir)
+        except OSError:
+            gem5_logger.warn("Failed to remove the temporary directory!")
+
+        # Delete the lock file
+        os.remove(self.lock_file_name)
+
+    # Functions only to be called by the Gem5 connection itself
+    def _connect_gem5_platform(self, platform):
+        port = platform.gem5_port
+        gem5_simulation = platform.gem5
+        gem5_interact_dir = platform.gem5_interact_dir
+        gem5_out_dir = platform.gem5_out_dir
+
+        self.connect_gem5(port, gem5_simulation, gem5_interact_dir, gem5_out_dir)
+
+    # This function connects to the gem5 simulation
+    def connect_gem5(self, port, gem5_simulation, gem5_interact_dir,
+                      gem5_out_dir):
+        """
+        Connect to the telnet port of the gem5 simulation.
+
+        We connect, and wait for the prompt to be found. We do not use a timeout
+        for this, and wait for the prompt in a while loop as the gem5 simulation
+        can take many hours to reach a prompt when booting the system. We also
+        inject some newlines periodically to try and force gem5 to show a
+        prompt. Once the prompt has been found, we replace it with a unique
+        prompt to ensure that we are able to match it properly. We also disable
+        the echo as this simplifies parsing the output when executing commands
+        on the device.
+        """
+        host = socket.gethostname()
+        gem5_logger.info("Connecting to the gem5 simulation on port {}".format(port))
+
+        # Check if there is no on-going connection yet
+        lock_file_name = '{}{}_{}.LOCK'.format(self.lock_directory, host, port)
+        if os.path.isfile(lock_file_name):
+            # There is already a connection to this gem5 simulation
+            raise TargetError('There is already a connection to the gem5 '
+                              'simulation using port {} on {}!'
+                              .format(port, host))
+
+        # Connect to the gem5 telnet port. Use a short timeout here.
+        attempts = 0
+        while attempts < 10:
+            attempts += 1
+            try:
+                self.conn = TelnetPxssh(original_prompt=None)
+                self.conn.login(host, self.username, port=port,
+                                login_timeout=10, auto_prompt_reset=False)
+                break
+            except pxssh.ExceptionPxssh:
+                pass
+        else:
+            gem5_simulation.kill()
+            raise TargetError("Failed to connect to the gem5 telnet session.")
+
+        gem5_logger.info("Connected! Waiting for prompt...")
+
+        # Create the lock file
+        self.lock_file_name = lock_file_name
+        open(self.lock_file_name, 'w').close() # Similar to touch
+        gem5_logger.info("Created lock file {} to prevent reconnecting to "
+                         "same simulation".format(self.lock_file_name))
+
+        # We need to find the prompt. It might be different if we are resuming
+        # from a checkpoint. Therefore, we test multiple options here.
+        prompt_found = False
+        while not prompt_found:
+            try:
+                self._login_to_device()
+            except TIMEOUT:
+                pass
+            try:
+                # Try and force a prompt to be shown
+                self.conn.send('\n')
+                self.conn.expect([r'# ', self.conn.UNIQUE_PROMPT, r'\[PEXPECT\][\\\$\#]+ '], timeout=60)
+                prompt_found = True
+            except TIMEOUT:
+                pass
+
+        gem5_logger.info("Successfully logged in")
+        gem5_logger.info("Setting unique prompt...")
+
+        self.conn.set_unique_prompt()
+        self.conn.prompt()
+        gem5_logger.info("Prompt found and replaced with a unique string")
+
+        # We check that the prompt is what we think it should be. If not, we
+        # need to update the regex we use to match.
+        self._find_prompt()
+
+        self.conn.setecho(False)
+        self._sync_gem5_shell()
+
+        # Fully connected to gem5 simulation
+        self.gem5_interact_dir = gem5_interact_dir
+        self.gem5_out_dir = gem5_out_dir
+        self.gem5simulation = gem5_simulation
+
+        # Ready for interaction now
+        self.ready = True
+
+    def _login_to_device(self):
+        """
+        Login to device, will be overwritten if there is an actual login
+        """
+        pass
+
+    def _find_prompt(self):
+        prompt = r'\[PEXPECT\][\\\$\#]+ '
+        synced = False
+        while not synced:
+            self.conn.send('\n')
+            i = self.conn.expect([prompt, self.conn.UNIQUE_PROMPT, r'[\$\#] '], timeout=self.default_timeout)
+            if i == 0:
+                synced = True
+            elif i == 1:
+                prompt = self.conn.UNIQUE_PROMPT
+                synced = True
+            else:
+                prompt = re.sub(r'\$', r'\\\$', self.conn.before.strip() + self.conn.after.strip())
+                prompt = re.sub(r'\#', r'\\\#', prompt)
+                prompt = re.sub(r'\[', r'\[', prompt)
+                prompt = re.sub(r'\]', r'\]', prompt)
+
+        self.conn.PROMPT = prompt
+
+    def _sync_gem5_shell(self):
+        """
+        Synchronise with the gem5 shell.
+
+        Write some unique text to the gem5 device to allow us to synchronise
+        with the shell output. We actually get two prompts so we need to match
+        both of these.
+        """
+        gem5_logger.debug("Sending Sync")
+        self.conn.send("echo \*\*sync\*\*\n")
+        self.conn.expect(r"\*\*sync\*\*", timeout=self.default_timeout)
+        self.conn.expect([self.conn.UNIQUE_PROMPT, self.conn.PROMPT], timeout=self.default_timeout)
+        self.conn.expect([self.conn.UNIQUE_PROMPT, self.conn.PROMPT], timeout=self.default_timeout)
+
+    def _gem5_util(self, command):
+        """ Execute a gem5 utility command using the m5 binary on the device """
+        if self.m5_path is None:
+            raise TargetError('Path to m5 binary on simulated system  is not set!')
+        self._gem5_shell('{} {}'.format(self.m5_path, command))
+
+    def _gem5_shell(self, command, as_root=False, timeout=None, check_exit_code=True, sync=True):  # pylint: disable=R0912
+        """
+        Execute a command in the gem5 shell
+
+        This wraps the telnet connection to gem5 and processes the raw output.
+
+        This method waits for the shell to return, and then will try and
+        separate the output from the command from the command itself. If this
+        fails, warn, but continue with the potentially wrong output.
+
+        The exit code is also checked by default, and non-zero exit codes will
+        raise a TargetError.
+        """
+        if sync:
+            self._sync_gem5_shell()
+
+        gem5_logger.debug("gem5_shell command: {}".format(command))
+
+        # Send the actual command
+        self.conn.send("{}\n".format(command))
+
+        # Wait for the response. We just sit here and wait for the prompt to
+        # appear, as gem5 might take a long time to provide the output. This
+        # avoids timeout issues.
+        command_index = -1
+        while command_index == -1:
+            if self.conn.prompt():
+                output = re.sub(r' \r([^\n])', r'\1', self.conn.before)
+                output = re.sub(r'[\b]', r'', output)
+                # Deal with line wrapping
+                output = re.sub(r'[\r].+?<', r'', output)
+                command_index = output.find(command)
+
+                # If we have -1, then we cannot match the command, but the
+                # prompt has returned. Hence, we have a bit of an issue. We
+                # warn, and return the whole output.
+                if command_index == -1:
+                    gem5_logger.warn("gem5_shell: Unable to match command in "
+                                     "command output. Expect parsing errors!")
+                    command_index = 0
+
+        output = output[command_index + len(command):].strip()
+
+        # It is possible that gem5 will echo the command. Therefore, we need to
+        # remove that too!
+        command_index = output.find(command)
+        if command_index != -1:
+            output = output[command_index + len(command):].strip()
+
+        gem5_logger.debug("gem5_shell output: {}".format(output))
+
+        # We get a second prompt. Hence, we need to eat one to make sure that we
+        # stay in sync. If we do not do this, we risk getting out of sync for
+        # slower simulations.
+        self.conn.expect([self.conn.UNIQUE_PROMPT, self.conn.PROMPT], timeout=self.default_timeout)
+
+        if check_exit_code:
+            exit_code_text = self._gem5_shell('echo $?', as_root=as_root,
+                                             timeout=timeout, check_exit_code=False,
+                                             sync=False)
+            try:
+                exit_code = int(exit_code_text.split()[0])
+                if exit_code:
+                    message = 'Got exit code {}\nfrom: {}\nOUTPUT: {}'
+                    raise TargetError(message.format(exit_code, command, output))
+            except (ValueError, IndexError):
+                gem5_logger.warning('Could not get exit code for "{}",\ngot: "{}"'.format(command, exit_code_text))
+
+        return output
+
+    def _mount_virtio(self):
+        """
+        Mount the VirtIO device in the simulated system.
+        """
+        gem5_logger.info("Mounting VirtIO device in simulated system")
+
+        self._gem5_shell('su -c "mkdir -p {}" root'.format(self.gem5_input_dir))
+        mount_command = "mount -t 9p -o trans=virtio,version=9p2000.L,aname={} gem5 {}".format(self.gem5_interact_dir, self.gem5_input_dir)
+        self._gem5_shell(mount_command)
+
+    def _move_to_temp_dir(self, source):
+        """
+        Move a file to the temporary directory on the host for copying to the
+        gem5 device
+        """
+        command = "cp {} {}".format(source, self.gem5_interact_dir)
+        gem5_logger.debug("Local copy command: {}".format(command))
+        subprocess.call(command.split())
+        subprocess.call("sync".split())
+
+    def _check_ready(self):
+        """
+        Check if the gem5 platform is ready
+        """
+        if not self.ready:
+            raise TargetError('Gem5 is not ready to interact yet')
+
+    def _wait_for_boot(self):
+        pass
+
+    def _probe_file(self, filepath):
+        """
+        Internal method to check if the target has a certain file
+        """
+        command = 'if [ -e \'{}\' ]; then echo 1; else echo 0; fi'
+        output = self.execute(command.format(filepath), as_root=self.is_rooted)
+        return boolean(output.strip())
+
+
+class LinuxGem5Connection(Gem5Connection):
+
+    def _login_to_device(self):
+        gem5_logger.info("Trying to log in to gem5 device")
+        login_prompt = ['login:', 'AEL login:', 'username:', 'aarch64-gem5 login:']
+        login_password_prompt = ['password:']
+        # Wait for the login prompt
+        prompt = login_prompt + [self.conn.UNIQUE_PROMPT]
+        i = self.conn.expect(prompt, timeout=10)
+        # Check if we are already at a prompt, or if we need to log in.
+        if i < len(prompt) - 1:
+            self.conn.sendline("{}".format(self.username))
+            password_prompt = login_password_prompt + [r'# ', self.conn.UNIQUE_PROMPT]
+            j = self.conn.expect(password_prompt, timeout=self.default_timeout)
+            if j < len(password_prompt) - 2:
+                self.conn.sendline("{}".format(self.password))
+                self.conn.expect([r'# ', self.conn.UNIQUE_PROMPT], timeout=self.default_timeout)
+
+
+
+class AndroidGem5Connection(Gem5Connection):
+
+    def _wait_for_boot(self):
+        """
+        Wait for the system to boot
+
+        We monitor the sys.boot_completed and service.bootanim.exit system
+        properties to determine when the system has finished booting. In the
+        event that we cannot coerce the result of service.bootanim.exit to an
+        integer, we assume that the boot animation was disabled and do not wait
+        for it to finish.
+
+        """
+        gem5_logger.info("Waiting for Android to boot...")
+        while True:
+            booted = False
+            anim_finished = True  # Assume boot animation was disabled on except
+            try:
+                booted = (int('0' + self._gem5_shell('getprop sys.boot_completed', check_exit_code=False).strip()) == 1)
+                anim_finished = (int(self._gem5_shell('getprop service.bootanim.exit', check_exit_code=False).strip()) == 1)
+            except ValueError:
+                pass
+            if booted and anim_finished:
+                break
+            time.sleep(60)
+
+        gem5_logger.info("Android booted")
+
 def _give_password(password, command):
     if not sshpass:
         raise HostError('Must have sshpass installed on the host in order to use password-based auth.')