#!/usr/bin/env python3

# Copyright (c) 2019 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

import os
import signal
import stat
import subprocess
import sys
import time

import CEosHelper
import TpmGeneric.Defs as TpmDefs
from TpmGeneric.Tpm import TpmGeneric

RC_EOS_PATH = '/mnt/flash/rc.eos'
SWI_RC_EOS_PATH = '/export/swi/rc.eos'

def execRc( rcfile ):
   p = None

   try:
      # Make sure the file is executable. On ext4 there is no guarantee that it
      # will be.
      st = os.stat( rcfile )
      if not st.st_mode & stat.S_IEXEC:
         os.chmod( rcfile, st.st_mode | stat.S_IEXEC )

      # We use the subprocess module instead of ManagedSubprocess
      # because we don't want to kill backgrounded processes upon
      # rc.eos's exit. ManagedSubprocess terminates background
      # processes when rc.eos exits. We run rc.eos before setting the
      # controlling terminal. If we run rc.eos after the setsid()
      # call, the background processes still get killed when rc.eos
      # exits.
      p = subprocess.Popen( [ rcfile ] ) # pylint: disable=consider-using-with

      # Install a signal handler for SIGINT, so that in case the script
      # hangs we can kill it.
      def handler( signum, frame ):
         # pylint: disable-next=consider-using-f-string
         print( 'Terminating %s with ^C' % rcfile, file=sys.stderr )
         sys.stderr.flush()
         # On some runs of EosInitHang, we may get the above output eaten by some
         # other prints on the console. As soon as we return, systemd starts other
         # services. Wait a little bit to allow time for the console to show the
         # message.
         time.sleep( 2 )
         sys.exit( 1 )

      signal.signal( signal.SIGINT, handler )

      tty = None
      # Find the console and make it the controlling terminal.
      with open( '/proc/cmdline' ) as f:
         cmdline = f.read()
      try:
         tty = ( ( cmdline[ cmdline.rindex( 'console=' ) + len( 'console=' ) : ] )
                    .partition( ' ' ) )[ 0 ].partition( ',' )[ 0 ]
      except ValueError as e: # pylint: disable=unused-variable
         pass

      # Inside a container, we still see the console info from outside the container.
      # Don't try to use it.
      if tty and not CEosHelper.isCeos():
         if '/dev/' in tty:
            tty_name = tty
         else:
            tty_name = '/dev/' + tty

         if 'uart8250' in tty_name:
            tty_name = '/dev/ttyS0'

         # Make ourselves session leader and open the tty. That should
         # make tty our controlling terminal.
         os.setsid()
         fd = os.open( tty_name, os.O_RDWR )

         # We want to listen to ^C. So dup stdinput to fd. We leave
         # stdoutput alone. Should we dup that too?
         os.dup2( fd, 0 )
         os.dup2( fd, 2 )

      p.wait()

   except OSError as e:
      print( f'Unexpected error when trying to access {rcfile}: {e}',
             file=sys.stderr )
      if p:
         p.kill()
         p.wait()

def maybeExecRc( rcFile=RC_EOS_PATH ):
   hasRcEos = os.path.exists( rcFile )

   # The extension of TPM PCRs is a security feature that allows post boot
   # attestation. We precisely measure every step of EOS that can load code. We
   # record hashes that can be extracted and validated by an external controller
   # in a secure manner. This can used as part of a security model to verify that
   # no tampering of any kind has happened on a device.
   #
   # rc.eos being one way to load a script very early at boot time, it's critical
   # to have a record of this execution.
   #
   # This is what the piece of code does below, if we detect a TPM, support the
   # feature on the platform, and it's enabled (via the 'measuredboot' bit in the
   # secure boot toggle)
   try:
      tpm = TpmGeneric()
      if ( tpm.isToggleBitSet( TpmDefs.SBToggleBit.MEASUREDBOOT ) and
           not rcFile.startswith( '/export/swi' ) ):
         if hasRcEos:
            tpm.pcrExtendFromFile( TpmDefs.PcrRegisters.EOS_EARLY_BOOT_CONF,
                                   TpmDefs.PcrEventType.EOS_RC_EOS,
                                   # pylint: disable-next=consider-using-f-string
                                   rcFile, log=[ '%s file found' % rcFile ] )
         else:
            tpm.pcrExtendFromData( TpmDefs.PcrRegisters.EOS_EARLY_BOOT_CONF,
                                   TpmDefs.PcrEventType.EOS_RC_EOS,
                                   TpmDefs.DefaultPcrExtend.NO_RC_EOS.value,
                                   # pylint: disable-next=consider-using-f-string
                                   log=[ 'no %s file found' % rcFile ] )
   except ( TpmDefs.NoTpmDevice, TpmDefs.NoTpmImpl, TpmDefs.NoSBToggle ):
      pass
   except TpmDefs.Error as e:
      # pylint: disable-next=consider-using-f-string
      print( 'Failed to extend PCR for rc.eos: %s' % str( e ) )

   if not hasRcEos:
      return

   print( rcFile, 'detected' )

   try:
      # We want to control the TTY to be able to CTRL+C rc.eos if it hangs. By
      # default systemd units cannot do that. Doing a fork() solves the problem.
      # That is not very elegant, but ...
      pid = os.fork()

      if pid == 0:
         execRc( rcFile )
      else:
         os.waitpid( pid, 0 )
   except OSError as e:
      print( 'Unexpected error while setting up rc.eos execution:', e,
             file=sys.stderr )

if __name__ == "__main__":
   maybeExecRc( RC_EOS_PATH )
   maybeExecRc( SWI_RC_EOS_PATH )
