#!/usr/bin/env python3
# Arista Networks, Inc. Confidential and Proprietary.
# Copyright (c) 2014 Arista Networks, Inc.  All rights reserved.
"""
1. Module implements a basic client-server file transfer with server 
   accepting only one connection. The interaction happens through files whose name
   is passed as an argument. 
2. All the calls to send/receive are blocking as this module is spawned as a 
   separate process.
3. Socket options are set to match MLAG TCP outbound socket connection which is
   used for attrLog updates.
"""
from __future__ import absolute_import, division, print_function
import six
import socket
import argparse 
import Tac
import MlagShared
import os
from Arnet.NsLib import NamespaceType, setns, DEFAULT_NS

# pylint: disable-msg=E1101
# ignore the complaint about no recv member in socket object

def isIpv6Addr( addr ):
   Af = Tac.Type( "Arnet::AddressFamily" )
   IpGenAddr = Tac.Type( 'Arnet::IpGenAddr' )
   return IpGenAddr( addr ).af == Af.ipv6

def setSocketOptions( sock, deviceName=None, ipv6=False ):
   for ( level, opt, val ) in \
         MlagShared.socketOpts( synCntOption=False, deviceName=deviceName,
                                ipv6=ipv6 ):
      if isinstance( val, str ):
         val = six.ensure_binary( val )
      sock.setsockopt( level, opt, val )

def runServer( ipAddr, port, devName, fileName ):
   blockSize = 1024 # number of bytes to block on receive call
   server = None
   print( "Running server process" )
   print( f"ipAddr: { ipAddr }, port: { port }, devName: { devName }" )
   try:
      if isIpv6Addr( ipAddr ):
         server = socket.socket( socket.AF_INET6, socket.SOCK_STREAM )
      else:
         server = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
      setSocketOptions( server, devName, ipv6=isIpv6Addr( ipAddr ) )
      server.bind( ( ipAddr, port ) )
   except socket.error as err:
      ( value, message ) = err.args # pylint: disable=unbalanced-tuple-unpacking
      if server:
         server.close()
      print( message )
      return value
   
   print( server.getsockname() )
   server.listen( 1 )
   conn, addr = server.accept()
   print( addr )
   setSocketOptions( conn, devName, ipv6=isIpv6Addr( ipAddr ) )
   with open( fileName, "wb" ) as fileStream:
      while True:
         data = conn.recv( blockSize )
         if not data:
            break
         fileStream.write( data )
   conn.close()
   server.close()
   print( "Exiting server process" )
   return 0

def runClient( ipAddr, port, devName, fileName ):
   client = None
   try:
      if isIpv6Addr( ipAddr ):
         client = socket.socket( socket.AF_INET6, socket.SOCK_STREAM )
      else:
         client = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
      setSocketOptions( client, devName, isIpv6Addr( ipAddr ) )
      client.connect( ( ipAddr, port ) )
   except socket.error as err:
      ( value, message ) = err.args # pylint: disable=unbalanced-tuple-unpacking
      if client:
         client.close()
      print( message )
      return value
   # Return value on successful file transfer
   result = 0
   print( "connected to server" )
   if os.path.isfile( fileName ):
      with open( fileName, "rb" ) as fileStream:
         client.sendall( fileStream.read() )
   else:
      # pylint: disable-next=consider-using-f-string
      print( "Non-existent file %s. IOError" % fileName )
      result = -1
   client.close()
   return result

# Returns client/server returnCode based on socket operation
def main():
   parser = argparse.ArgumentParser()
   group = parser.add_mutually_exclusive_group()
   group.add_argument( '--client', action="store_true", default=True,
                        help="Client process" )
   group.add_argument( '--server', action="store_true",
                        help="Server process" )
   parser.add_argument( '-ip', '--ipAddr', action="store", default='127.0.0.1',
                        help="Ip address for client/server" )
   parser.add_argument( '--port', type=int, action="store", default=10000,
                        help="TCP port for client/server" )
   parser.add_argument( '-dev', '--device', action="store", default='lo',
                        help="Device name to bind the socket" )
   parser.add_argument( '-f', '--fileName', action='store', 
                        default='/tmp/testFile.txt', 
                        help="File name for read/write for client/server" )
   parser.add_argument( '-n', '--nameSpace', action='store',
                        default=DEFAULT_NS,
                        help="namespace for client/server" )
   args = parser.parse_args()
   result = -1

   if args.nameSpace != DEFAULT_NS:
      # change to the local-interface namespace for the entirety of the process
      # pylint: disable-next=consider-using-with,consider-using-f-string
      ns = open( '/var/run/netns/%s' % args.nameSpace )
      setns( ns.fileno(), nsType=NamespaceType.NETWORK )
   if args.server:
      result = runServer( args.ipAddr, args.port, args.device, args.fileName )
   elif args.client:
      result = runClient( args.ipAddr, args.port, args.device, args.fileName )
   else:
      assert False, "Client/Server are the only two args. Invalid input."
   
   print( "result %d" % result ) # pylint: disable=consider-using-f-string
   return result

if __name__ == '__main__':
   exit( main() ) # pylint: disable=consider-using-sys-exit
