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

import Tac, sys
import Ethxmit
import ctypes
import optparse
from socket import socket, AF_PACKET, SOCK_RAW

# txraw header used for passsing information needed for selecting
# chip for inseting HiGig2 frame for L2 unicast learning

class StrataTxRaw():
   def __init__( self ):
      self.txraw_bs = bytearray( 4 )
      
   def setIngressModId( self, modId ):
      self.txraw_bs[ 0 ] = modId

   def getByteArray( self ):
      return self.txraw_bs

# HiGig2 is properietary header used by Broadcom for links running in HiGig2 mode

class HiGig2():
   def __init__( self ):
      self.opcodeMap = { "unicast" : 1, "l2mc" : 3, "ipmc" : 4 }
      self.hg2_bs = bytearray( 16 )
      self.hg2_bs[ 0 ] = 0xFC # start

   def setSrcModId( self, modId):
      self.hg2_bs[ 4 ] = modId

   def setSrcPortId( self, portId):
      self.hg2_bs[ 5 ] = portId

   def setDstPortId( self, portId):
      self.hg2_bs[ 3 ] = portId

   def setDstModId( self, modId):
      self.hg2_bs[ 2 ] = modId

   def setMcGid( self, mcGroup ):
      self.hg2_bs[ 2 ] = ( mcGroup >> 8 ) & 0xFF
      self.hg2_bs[ 3 ] = ( mcGroup & 0xFF )

   def setOpCode ( self, opCode ):
      self.hg2_bs[ 14 ] = self.opcodeMap[ opCode ]

   def setVlan ( self, vlan ):
      self.hg2_bs[ 13 ] = ( vlan & 0xFF )
      self.hg2_bs[ 12 ] = ( vlan >> 8 ) & 0x0F

   def setCongestionClass( self, cc ):
      self.hg2_bs[ 7 ] = ( self.hg2_bs[ 7 ] & 0xE7 ) | ( ( cc & 0x3 ) << 3 )

   def getByteArray( self ):
      return self.hg2_bs

parser = optparse.OptionParser()
parser.usage = "%prog [ OPTIONS ]"

parser.add_option( "-t", "--txraw", type="int", dest="rawHdrType", default=0,
                   help="txraw header, 0 - Strata (default: 0)" )
parser.add_option( "-i", "--ingress-mod-id", type="int", dest="ingressModId",
                   default=0, help="HiGig2 source module id (default: 0)" )
parser.add_option( "-m", "--src-mod-id", type="int", dest="srcModId", default=0,
                   help="HiGig2 source module id (default: 0)" )
parser.add_option( "-p", "--src-port-id", type="int", dest="srcPortId", default=0,
                   help="HiGig2 source port id (default: 0)" )
parser.add_option( "-M", "--dst-mod-id", type="int", dest="dstModId", default=0,
                   help="HiGig2 dest module id (default: 0)" )
parser.add_option( "-P", "--dst-port-id", type="int", dest="dstPortId", default=0,
                   help="HiGig2 dest port id (default: 0)" )
parser.add_option( "-d", "--mcg-id", type="int", dest="mcGroupId", default=0,
                   help="HiGig2 multicast group id (default: 0)" )
parser.add_option( "-o", "--opcode", type="string", dest="opCode", default="l2mc",
                   help="HiGig2 opcode: unicast, l2mc, ipmc  (default: l2mc )" )
parser.add_option( "-v", "--vlan", type="int", dest="vlan", default=1,
                   help="HiGig2 vlan (default: 1)" )
parser.add_option( "--cc", type="int", dest="cc", default=0,
                   help="HiGig2 congestion class" )
parser.add_option( "-s", "--size", type="int", dest="size", default=0,
                   help="payload size (default: 0)" )

# Options controlling how many packets are transmitted.
parser.add_option( "-n", "--num", type="int", default=1,
                   help="total number of packets to send (default: %default)" )
parser.add_option( "-b", "--burst", type="int", default=1,
                   help="number of packets in each burst (default: %default)" )
parser.add_option( "--sleep", type="float", default=0.0,
                   help="how long to sleep (in seconds) after each burst "
                   "(default: %default)")

# Options controlling packet contents
parser.add_option( "--ip", action="store_true", dest="ip",
                   help="payload becomes an IP packet with some default values" )

( options, args ) = parser.parse_args()

if len( args ) == 0:
   pass
else:
   parser.error( "unexpected args: %s" % args[ 1: ] )
   sys.exit( 0 )

txIntf = "txraw"

if options.size:
   payload = bytearray( options.size )
   for i in range( 0, options.size ):
      payload[ i ] = i % 255

if options.ip:
   size = 100
   sip = '10.1.1.1'
   dip = '10.1.1.2'
   tos = options.cc & 0x3 # tos = ([dscp=0] << 2) | [ECN=CC]
   ipPacket = Ethxmit.buildIpPacket( size, sip, dip, tos=options.cc,
                                     extraHeaderSize=4 ) # 4 bytes for Dot1q
   smac = '00:11:22:33:44:55'
   dmac = '00:11:22:33:44:56'
   ethertype = 0x0800 # IPv4
   ethHeader = Ethxmit.buildEthernetHeader( dmac, smac, ethertype,
                                            vlan=options.vlan )
   payload = ethHeader + ipPacket

if options.rawHdrType is 0:

   # prepare txraw header
   txraw = StrataTxRaw( )
   txraw.setIngressModId( options.ingressModId )

   # prepare higig2 header using arguments passed in
   hg2 = HiGig2()
   hg2.setSrcModId( options.srcModId )
   hg2.setSrcPortId( options.srcPortId )
   if options.opCode == 'unicast':
      hg2.setDstModId( options.dstModId )
      hg2.setDstPortId( options.dstPortId )
   else:
      hg2.setMcGid( options.mcGroupId )
   hg2.setCongestionClass( options.cc )   
   hg2.setOpCode( options.opCode )
   hg2.setVlan( options.vlan )
   parray = txraw.getByteArray() + hg2.getByteArray() + payload
   
   # Prepare packet trailer here FCS + EOP
   parray += bytearray( 5 )
   parray[ len( parray ) - 1 ] = 0xFD

   # PAD with IDLE to make it aligned at 4 bytes
   pad = len( parray ) % 4
   if pad:
      align = 4 - pad
      parray += bytearray( align )
      for i in range( 1, align + 1 ):
         parray[ len( parray ) - i ] = 0x07
else:
   parser.error( "unsupported txraw header: %s" % options.rawHdrType )
   sys.exit( 0 )

# Open raw packet socket
sock = socket( AF_PACKET, SOCK_RAW )
sock.bind( ( txIntf, 0 ) )

et = ctypes.cdll.LoadLibrary( "libethxmit.so" )
et.socksend( sock.fileno(), ctypes.c_char_p( str( parray ) ), len( parray ),
             options.num, options.burst, ctypes.c_float( options.sleep ),
             0, 0, 0, 0 )
