#!/usr/bin/env python
# Copyright (c) 2015 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

import errno
import os
import select
import signal
import socket
import sys
import threading
import time

import FastServ
import FastServUtil

ONE_TIME_LAUNCH_SERVER_ADDRESS_FMT = \
  "\0/fastclient/oneTimeLaunch/server%s" # pylint: disable-msg=W1401

# munch our argvs
appBinary = sys.argv[1] # [0] is 'this program' (/usr/bin/oneTimeLaunchBackend)
serverAddr = sys.argv[2] # address we should listen for frontend connection from
argList = sys.argv[3:]
argList.insert( 0, appBinary )

def createServerSocket( address ):
   sock_ = socket.socket( socket.AF_UNIX, socket.SOCK_STREAM, 0 )
   sock_.bind( address )
   sock_.listen( 1 )
   sock_.setsockopt( socket.SOL_SOCKET, socket.SO_REUSEADDR, 1 )
   sock_.setblocking( 0 )
   return sock_

childPid = None # the app we will start once we have gotten IO from frontend
signalSocket = None # we get that one once we have a connection from frontend
serverSocket = createServerSocket( ONE_TIME_LAUNCH_SERVER_ADDRESS_FMT % serverAddr )
socketsToRead = [ serverSocket ]

def waitChild( pid, signalSocket_ ):
   ( pid, status ) = os.waitpid( pid, 0 )
   #print "child pid", pid, "exited, status", status
   # write the return code back to the frontend and go away
   FastServUtil.writeInteger( signalSocket_, status ) 
   signalSocket_.close()
   os._exit( 0 ) # pylint: disable-msg=W0212

# pylint: disable-next=inconsistent-return-statements
def spawnChild( appBinary_, argList_ ):
   pid = os.fork()
   if pid > 0:
      return pid
   os.execvp( appBinary_, argList_ )

def handleSignals( signalSocket_, childPid_ ):
   """ A signal was recieved on the frontend and forwarded here via fastclient's
       signal socket: propagate it to the child (the actuall app). """
   signalNumber = FastServUtil.readInteger( signalSocket_ )
   if not signalNumber:
      # Frontend is gone (socket closed), kill the app since its controlling
      # terminal gone it has not more reasons to live (we don't support nohup).
      os.killpg( childPid_, signal.SIGHUP )
      os.killpg( childPid_, signal.SIGCONT )
      # This process may not sleep this long. If SIGHUP is able to terminates child,
      # the 'waitChild' thread will catch its exit code and exit itself, otherwise
      # the big hammer of the ungracefull SIGKILL will fall in 5 seconds.
      time.sleep( 5 ) 
      os.killpg( childPid_, signal.SIGKILL )
      
   else:
      #print "forwarding signal", signalNumber
      if signalNumber < 0:
         kpid = -childPid_
         ksig = -signalNumber # pylint: disable=invalid-unary-operand-type
      else:
         kpid = childPid_
         ksig = signalNumber
         
      os.kill( kpid, ksig ) #TODO: Do we need to have a try catch here?

# main loop: wait for connection, setup IO, start app, then wait for signal socket 
# to bring in some signals to be forwarded or for the signal socket to be closed. 
# If something nefarious happened to the frontend, then ProcMgr will kill us when 
# the procmgr config file is removed (so infinite wait on the select is fine).
while( True ): # pylint: disable=superfluous-parens
   try:
      socks, _, _ = select.select( socketsToRead, [], [] )
   except OSError as e:
      if e.args[ 0 ] == errno.EINTR:
         continue
      raise # if we get interrupted by something else we should blow up!
   except KeyboardInterrupt:
      os._exit( 0 ) # pylint: disable-msg=W0212
   
   for sock in socks:
      if sock is signalSocket:
         handleSignals( signalSocket, childPid )
         # may never return
      if sock is serverSocket:
         # this case will ever only match the first time (we close it down there)
         conn, _ = sock.accept() # conn.fileno() is socket
         serverSocket.close()
         # send back some dummy integer that some other usages of fastclient need.
         try:
            FastServUtil.writeInteger( conn, 0 )
         except OSError:
            os._exit( -2 ) # pylint: disable-msg=W0212
         # tell front-end to send us its terminal
         FastServUtil.writeString( conn, 'n' )
         # get the signal fd (over which we get signals from the client or that
         # we use to return the exit status of the application), that's the only
         # socket we read from now (we closed the server socket already), and the
         # connection will be closed shortly too after getting the terminal from it.
         # pylint: disable-next=c-extension-no-member
         signalFd = FastServ.recvFds( conn.fileno(), 1 )
         if not signalFd:
            os._exit( -3 ) # pylint: disable-msg=W0212
         signalSocket = socket.fromfd( signalFd[ 0 ], socket.AF_UNIX,
                                                      socket.SOCK_STREAM, 0 )
         socketsToRead = [ signalSocket ]

         # read the type of connection from the connection. If a terminal, get the
         # name of the terminal and open it, if the io is passed as 3 fds, collect
         # them.
         fds = None
         ttyType = conn.recv( 1 )
         if ttyType == 't':
            ttyName = FastServUtil.readString( conn )
            ttyFile = open( ttyName, 'a+' ) # pylint: disable=consider-using-with
            fileNo = os.dup( ttyFile.fileno() )
            fds = { "stdin": fileNo, "stdout": fileNo, "stderr": fileNo }
         elif ttyType == 'u':
            # pylint: disable-next=c-extension-no-member
            fds = FastServ.recvFds( conn.fileno(), 3 )
            fds = { "stdin": fds[ 0 ], "stdout": fds[ 1 ], "stderr": fds[ 2 ] }
         else:
            # pylint: disable-next=consider-using-f-string
            raise RuntimeError( "Unknown TTY type %s" % ttyType )
         # set our IO to what the client wants (collected just above)
         os.dup2( fds[ "stdin" ], 0 )
         os.dup2( fds[ "stdout" ], 1 )
         os.dup2( fds[ "stderr" ], 2 )
         sys.stdin = os.fdopen( 0, "r" )
         sys.stdout = os.fdopen( 1, "w" )
         sys.stderr = os.fdopen( 2, "w" )
         # now get the env vars from the client and insert them into our env 
         fastClientArgs = FastServUtil.processEnvArgs( conn.fileno() )
         _ = FastServUtil.readString( conn ) # eat some passive mount related stuff
         #sys.stdout.flush()
         # connection has given us all it had, close it
         conn.close()
         # start our application and start a watcher thread 
         childPid = spawnChild( appBinary, argList )
         # monitor the child in another thread, main thread  will continue sending
         # to the child if we recieve any from the frontend.
         threading.Thread( target=waitChild, args=(childPid, signalSocket) ).start()

