#!/usr/bin/env python3
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This software product is a proprietary product of Mellanox Technologies Ltd.
# (the "Company") and all right, title, and interest in and to the software
# product, including all associated intellectual property rights, are and
# shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.
#

import argparse
import sys
import shlex
import json
import socket
import time
import os
import copy

__version__='v25.07.2'

try:
    from shlex import quote
except ImportError:
    from pipes import quote

MODIFY_SUBCMD_DEVICE = 0
MODIFY_SUBCMD_QUEUE = 1
DEBUG_SUBCMD_MEM = 2
DEBUG_SUBCMD_RING = 3
DEBUG_SUBCMD_TRIGGER = 4
DEBUG_SUBCMD_OPTIONS = 5
DEBUG_SUBCMD_COREDUMP = 6
DEBUG_SUBCMD_QUERY = 7
DEBUG_SUBCMD_DEV_MEM = 8
DEBUG_SUBCMD_AARFS_HASH_DUMP = 9
DEBUG_SUBCMD_AARFS_MRU_DUMP = 10
DEBUG_SUBCMD_AARFS_STATS_DUMP = 11
DEBUG_SUBCMD_NETDIM_STATS = 12

HEALTH_SUBCMD_SHOW = 0
HEALTH_SUBCMD_RECOVER = 1

class JsonRpcSnapException(Exception):
    def __init__(self, message):
        self.message = message

class JsonRpcVirtnetClient(object):
    decoder = json.JSONDecoder()

    def __init__(self, address, port, timeout=60.0):
        self.sock = None
        self._request_id = 0
        self.timeout = timeout
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((address, int(port)))
        except socket.error as ex:
            print("ERR: Can't connect to virtnet controller: %s" % (ex))
            print("     Check 'systemctl status virtio-net-controller'")
            print("     Or controller is not ready to accept commands")
            sys.exit()

    def __json_to_string(self, request):
        return json.dumps(request)

    def send(self, method, params=None):
        self._request_id += 1
        req = {
            'jsonrpc': '2.0',
            'method': method,
            'id': self._request_id
        }
        if params:
            req['params'] = copy.deepcopy(params)

        self.sock.sendall(self.__json_to_string(req).encode("utf-8"))
        return self._request_id

    def __string_to_json(self, request_str):
        try:
            obj, idx = self.decoder.raw_decode(request_str)
            return obj
        except ValueError:
            return None

    def recv(self):
        start_time = time.process_time()
        response = None
        buf = ""

        while not response:
            try:
                timeout = self.timeout - (time.process_time() - start_time)
                if timeout < 0:
                    break
                self.sock.settimeout(timeout)
                buf += self.sock.recv(4096).decode("utf-8")
                response = self.__string_to_json(buf)
            except socket.timeout:
                break
            except ValueError:
                continue  # incomplete response; keep buffering

        self.sock.close()

        if not response:
            raise JsonRpcSnapException("Response Timeout")
        return response

    def call(self, method, params={}):
        if params:
            hex_params = {
                key:
                    hex(value) if isinstance(value, int)
                    else value
                for key, value in params.items()
            }
            print(hex_params)
        req_id = self.send(method, params)
        response = self.recv()

        if 'error' in response:
            params["method"] = method
            params["req_id"] = req_id
            msg = "\n".join(["request:", "%s" % json.dumps(params, indent=2),
                             "Got JSON-RPC error response",
                             "response:",
                             json.dumps(response['error'], indent=2)])
            raise JsonRpcSnapException(msg)

        return response['result']

def call_rpc_func(args):
    return args.func(args)

def execute_script(parser, client, fd):
    executed_rpc = ""
    for rpc_call in map(str.rstrip, fd):
        if not rpc_call.strip():
            continue
        executed_rpc = "\n".join([executed_rpc, rpc_call])
        args = parser.parse_args(shlex.split(rpc_call))
        args.client = client
        try:
            call_rpc_func(args)
        except JsonRpcSnapException as ex:
            print("Exception:")
            print(executed_rpc.strip() + " <<<")
            print(ex.message)
            exit(1)

def hotplug(args):
    params = {
        'ib_device': args.ib_device,
    }
    if args.num_queues == None and args.max_queues != None:
        args.num_queues = args.max_queues
    if (args.max_queue_pairs != None and args.num_queues != None) or \
        (args.max_queue_pairs == None and args.num_queues == None):
        print("Please specify either max_queue_pairs or num_queues")
        return -1
    if args.mac != None:
        params['mac'] = args.mac
    if args.mtu != None:
        params['mtu'] = args.mtu
    if args.num_queues != None:
        params['num_queues'] = args.num_queues
    if args.max_queue_size != None:
        params['max_queue_size'] = args.max_queue_size
    if args.sf_num != None:
        params['sf_num'] = args.sf_num
    if args.features != None:
        params['features'] = args.features
    if args.legacy != None:
        params['legacy'] = args.legacy
    if args.max_queue_pairs != None:
        if args.max_queue_pairs < 1:
            print("max_queue_pairs is at least 1 for hotplug")
            return -1
        params['max_queue_pairs'] = args.max_queue_pairs

    result = args.client.call('hotplug', params)
    print(json.dumps(result, indent=2))
    return result['errno']

def unplug(args):
    params = {}
    if (args.pf != None and args.vuid != None) or \
        (args.pf == None and args.vuid == None):
        print("Please specify either -p or -u")
        return -1
    if args.pf != None:
        params['id'] = args.pf
    if args.vuid != None:
        params['vuid'] = args.vuid
    result = args.client.call('unplug', params)
    print(json.dumps(result, indent=2))
    return result['errno']

def _list(args):
    result = args.client.call('list')
    print(json.dumps(result, indent=2))

def query(args):
    params = {}
    if args.all != None:
        params['all'] = args.all
    if (args.pf != None and args.vuid != None):
        print("Please specify either \"-p\" or \"-u\"")
        return -1
    if args.vf != None:
        if args.pf != None:
            params['vf'] = args.vf
        else:
            print("Please specify either \"-p + -v\" or \"-u\"")
            return -1
    if args.pf != None:
        params['pf'] = args.pf
    if args.vuid != None:
        params['vuid'] = args.vuid
    if args.dbg_stats != None:
        params['dbg_stats'] = args.dbg_stats
    if args.brief != None:
        params['brief'] = args.brief
    if args.latency_stats != None:
        params['latency_stats'] = args.latency_stats
    if args.queue_id != None:
        params['queue_id'] = args.queue_id
    if args.stats_clear != None:
        params['stats_clear'] = args.stats_clear

    result = args.client.call('query', params)
    print(json.dumps(result, indent=2))

def modify(args):
    params = {}
    if args.all == False:
        if (args.pf != None and args.vuid != None) or \
            (args.pf == None and args.vuid == None):
            print("Please specify either \"-p\" or \"-u\"")
            return -1
        if args.vf != None:
            if args.pf != None:
                params['vf'] = args.vf
            else:
                print("Please specify either \"-p + -v\" or \"-u\"")
                return -1
    else:
        if (args.pf != None or args.vuid != None or args.vf != None):
            print("Ignore device/queue related parameters")
    if args.pf != None:
        params['pf'] = args.pf
    if args.vuid != None:
        params['vuid'] = args.vuid
    if args.all != None:
        params['all'] = args.all
    if args.subcmd == 'device':
        params['subcmd'] = MODIFY_SUBCMD_DEVICE
        if (args.max_queue_pairs != None and args.msix_num != None):
            print("Please specify either max_queue_pairs or msix_num")
            return -1
        if args.state != None:
            params['state'] = args.state
        if args.features != None:
            params['features'] = args.features
        if args.mac != None:
            params['mac'] = args.mac
        if args.mtu != None:
            params['mtu'] = args.mtu
        if args.speed != None:
            params['speed'] = args.speed
        if args.link != None:
            params['link'] = args.link
        if args.rx_mode != None:
            params['rx_mode'] = args.rx_mode
        if args.msix_num != None:
            params['msix_num'] = args.msix_num
        if args.max_queue_size != None:
            params['max_queue_size'] = args.max_queue_size
        if args.supported_hash_types != None:
            params['supported_hash_types'] = args.supported_hash_types
        if args.rss_max_key_size != None:
            params['rss_max_key_size'] = args.rss_max_key_size
        if args.dst_port != None:
            params['dst_port'] = args.dst_port
        if args.rx_dma_q_num != None:
            if args.all == False:
                print("WARN: Applied to all devices, device/queue inputs are ignored");
                params['all'] = True
            params['rx_dma_q_num'] = args.rx_dma_q_num
        if args.rx_dim_config != None:
            params['rx_dim_config'] = args.rx_dim_config
        if args.drop_counter != None:
            params['drop_counter'] = args.drop_counter
        if args.packet_counter != None:
            params['packet_counter'] = args.packet_counter
        if args.aarfs_config != None:
            params['aarfs_config'] = args.aarfs_config
            if args.vf != None:
                print("ERROR: aarfs_config only applies to PF")
        if args.max_queue_pairs != None:
            params['max_queue_pairs'] = args.max_queue_pairs
        if args.dim_config != None:
            params['dim_config'] = args.dim_config
    elif args.subcmd == 'queue':
        params['subcmd'] = MODIFY_SUBCMD_QUEUE
        if args.period_mode != None:
            params['period_mode'] = args.period_mode
        if args.period != None:
            params['period'] = args.period
        if args.max_count != None:
            params['max_count'] = args.max_count

    result = args.client.call('modify', params)
    print(json.dumps(result, indent=2))
    return result['errno']

def log_level(args):
    params = {
        'level': args.level,
    }

    result = args.client.call('log_level', params)
    print(json.dumps(result, indent=2))

def validate(args):
    if args.file is None:
        print("Error: No input file provided.")
        return 1

    try:
        with open(args.file, 'r') as f:
            data = json.load(f)
            if isinstance(data, (dict, list)):
                print(args.file + ' is valid')
                return 0
            else:
                print(args.file + ' is not valid (not a JSON object or array)')
                return 1
    except (ValueError, json.JSONDecodeError) as e:
        print(f"Error: {e}")
        print("Try to format it on https://jsonformatter.curiousconcept.com")
        return 1
    except Exception as e:
        print(f"Unexpected error: {e}")
        return 1

def version(args):
    params = {}

    result = args.client.call('version', params)
    print(json.dumps(result, indent=2))

def restart(args):
    params = {}

    result = args.client.call('restart', params)
    print(json.dumps(result, indent=2))

def update(args):
    params = {}
    if (args.start == 1 and args.status == 1) or \
        (args.start == 0 and args.status == 0):
        print("Please specify either -s or -t")
        return -1
    if args.start == 1:
        params['start'] = args.start
    if args.status == 1:
        params['status'] = args.status

    result = args.client.call('update', params)
    print(json.dumps(result, indent=2))

def debug(args):
    params = {}
    if (args.pf != None and args.vuid != None):
        print("Please specify either \"-p\" or \"-u\"")
        return -1
    if args.vf != None:
        if args.pf != None:
            params['vf'] = args.vf
        else:
            print("Please specify either \"-p + -v\" or \"-u\"")
            return -1
    if args.pf != None:
        params['pf'] = args.pf
    if args.vuid != None:
        params['vuid'] = args.vuid
    if hasattr(args, 'subcmd'):
        if args.subcmd == 'mem':
            params['subcmd'] = DEBUG_SUBCMD_MEM
            if args.address != None:
                params['address'] = args.address
            if args.num_dword != None:
                params['num_dword'] = args.num_dword
        if args.subcmd == 'ring':
            params['subcmd'] = DEBUG_SUBCMD_RING
            if args.queue_id != None:
                params['queue_id'] = args.queue_id
            if args.ring_type != None:
                params['ring_type'] = args.ring_type
            if args.start != None:
                params['start'] = args.start
            if args.count != None:
                params['count'] = args.count
        if args.subcmd == 'trigger':
            params['subcmd'] = DEBUG_SUBCMD_TRIGGER
            if args.queue_id != None:
                params['queue_id'] = args.queue_id
        if args.subcmd == 'options':
            params['subcmd'] = DEBUG_SUBCMD_OPTIONS
            if args.dma_q_rec != None:
                params['dma_q_rec'] = args.dma_q_rec
            if args.lm_log != None:
                params['lm_log'] = args.lm_log
        if args.subcmd == 'coredump':
            params['subcmd'] = DEBUG_SUBCMD_COREDUMP
        if args.subcmd == 'query':
            params['subcmd'] = DEBUG_SUBCMD_QUERY
        if args.subcmd == 'dev_mem':
            params['subcmd'] = DEBUG_SUBCMD_DEV_MEM
            if args.address != None:
                params['address'] = args.address
            if args.num_dword != None:
                params['num_dword'] = args.num_dword
            if args.file != None:
                params['file'] = args.file
        if args.subcmd == 'dump_hash':
            params['subcmd'] = DEBUG_SUBCMD_AARFS_HASH_DUMP
            if args.pf != None:
                params['pf'] = args.pf
        if args.subcmd == 'dump_mru':
            params['subcmd'] = DEBUG_SUBCMD_AARFS_MRU_DUMP
        if args.subcmd == 'dump_stats':
            params['subcmd'] = DEBUG_SUBCMD_AARFS_STATS_DUMP
        if args.subcmd == 'netdim':
            params['subcmd'] = DEBUG_SUBCMD_NETDIM_STATS
            if args.queue_id != None:
                params['queue_id'] = args.queue_id
    else:
        params['subcmd'] = -1

    result = args.client.call('debug', params)
    print(json.dumps(result, indent=2))

def health(args):
    params = {}
    if (args.pf != None and args.vuid != None):
        print("Please specify either \"-p\" or \"-u\"")
        return -1
    if args.vf != None:
        if args.pf != None:
            params['vf'] = args.vf
        else:
            print("Please specify either \"-p + -v\" or \"-u\"")
            return -1
    if args.pf != None:
        params['pf'] = args.pf
    if args.vuid != None:
        params['vuid'] = args.vuid
    if args.all != None:
        params['all'] = args.all
    if hasattr(args, 'subcmd'):
        if args.subcmd == 'show':
            params['subcmd'] = HEALTH_SUBCMD_SHOW
    else:
        params['subcmd'] = -1

    result = args.client.call('health', params)
    print(json.dumps(result, indent=2))
def stats(args):
    params = {}
    if (args.pf != None and args.vuid != None) or \
        (args.pf == None and args.vuid == None):
        print("Please specify either -p or -u")
        return -1
    if args.vf != None:
        if args.pf != None:
            params['vf'] = args.vf
        else:
            print("Please specify either \"-p + -v\" or \"-u\"")
            return -1
    if args.pf != None:
        params['pf'] = args.pf
    if args.vuid != None:
        params['vuid'] = args.vuid
    if args.queue_id != None:
        params['queue_id'] = args.queue_id

    result = args.client.call('stats', params)
    print(json.dumps(result, indent=2))

def int_hex(x):
    return int(x, 16)

def main():
    manager_server_port='12191'
    health_server_port='22190'
    server_addr='127.0.0.1'
    server_port='12190'
    timeout=100.0

    parser = argparse.ArgumentParser(
        description='Nvidia virtio-net-controller command line interface ' + __version__)
    parser.add_argument('-v', '--version', action='version', version=__version__)
    subparsers = parser.add_subparsers(help='** Use -h for sub-command usage',
                                       dest='called_rpc_name')

    # Hotplug
    p = subparsers.add_parser('hotplug', help='hotplug virtnet device',
                              formatter_class=argparse.RawTextHelpFormatter)
    p.add_argument('-i', '--ib_device', help="IB device name for creating "
                   "data path SF, e.g, mlx5_0",
                   required=True, type=str)
    p.add_argument('-m', '--mac', help="MAC in format 00:11:22:33:44:55",
                   required=True, type=str)
    p.add_argument('-t', '--mtu', help="MTU", required=True, type=int)
    p.add_argument('-n', '--num_queues', help="max number of queues, at least 3",
                   required=False, type=int)
    p.add_argument('-q', '--max_queues', help=argparse.SUPPRESS, required=False,
                   type=int)
    p.add_argument('-s', '--max_queue_size', help="max depth of each queue",
                   required=True, type=int)
    p.add_argument('-u', '--sf_num', help="SF number", required=False, type=int)
    p.add_argument('-f', '--features', help="Feature bits in hex: "
                   "e.g, 0x80000 for vlan"
                   "\nVIRTIO_NET_F_CSUM             0"
                   "\nVIRTIO_NET_F_GSO              6"
                   "\nVIRTIO_NET_F_GUEST_TSO4       7"
                   "\nVIRTIO_NET_F_GUEST_TSO6       8"
                   "\nVIRTIO_NET_F_HOST_TSO4        11"
                   "\nVIRTIO_NET_F_HOST_TSO6        12"
                   "\nVIRTIO_NET_F_CTRL_VLAN        19"
                   "\nVIRTIO_NET_F_GUEST_ANNOUNCE   21"
                   "\nVIRTIO_NET_F_HASH_REPORT      57",
                   required=False, type=int_hex)
    p.add_argument('-l', '--legacy', help='Create legacy (transitional) device',
                   required=False, action="store_const", const=1, default=0)
    p.add_argument('-qp', '--max_queue_pairs', help="Modify maximum number of "
		    "queue pairs, each pair has 1 tx and 1 rx",
		    required=False, type=int)
    p.set_defaults(func=hotplug)

    # Unplug
    p = subparsers.add_parser('unplug', help='unplug virtnet device')
    p.add_argument('-p', '--pf', help="PF device ID", required=False, type=int)
    p.add_argument('-u', '--vuid', help="PF device VUID", required=False, type=str)
    p.set_defaults(func=unplug)

    # List
    p = subparsers.add_parser('list', help='list all virtnet devices')
    p.set_defaults(func=_list)

    # Query
    p = subparsers.add_parser('query', help='query all or individual virtnet '
                              'device(s)')
    p.add_argument('-a', '--all', help='query all virtnet devices, including '
                    'both PFs and VFs', required=False, action="store_const",
                    const=1, default=0)
    p.add_argument('-p', '--pf', help="PF ID, refer to list command",
                   required=False, type=int)
    p.add_argument('-v', '--vf', help="VF ID, refer to list command",
                   required=False, type=int)
    p.add_argument('-u', '--vuid', help="PF/VF device VUID", required=False, type=str)
    p.add_argument('--dbg_stats', help="Print debug counters and information",
                   required=False, action="store_true", default=False)
    p.add_argument('-b', '--brief', help="query brief info only",
                   required=False, action="store_true", default=False)
    p.add_argument('--latency_stats', help="Print latency cycles and information",
                   required=False, action="store_true", default=False)
    p.add_argument('-q', '--queue_id', help="Queue index of the device",
                   required=False, type=int)
    p.add_argument('--stats_clear', help="Clear stats",
                   required=False, action="store_true", default=False)
    p.set_defaults(func=query)

    # Modify
    p = subparsers.add_parser('modify', help='modify virtnet device')
    p.add_argument('-p', '--pf', help="pf ID", required=False, type=int)
    p.add_argument('-v', '--vf', help="vf ID", required=False, type=int)
    p.add_argument('-u', '--vuid', help="PF device VUID", required=False, type=str)
    p.add_argument('-a', '--all', help='Modify all virtnet devices/queues',
                   required=False, action="store_true", default=False)
    sp = p.add_subparsers()
    dp = sp.add_parser('device', help='modify device related options',
                       formatter_class=argparse.RawTextHelpFormatter);
    dp.add_argument('-m', '--mac', help="MAC in format 00:11:22:33:44:55",
                   required=False, type=str)
    dp.add_argument('-t', '--mtu', help="MTU", required=False, type=int)
    dp.add_argument('-e', '--speed', help="Speed in Mbps: "
                   "\n1000"
                   "\n2500"
                   "\n5000"
                   "\n10000"
                   "\n20000"
                   "\n25000"
                   "\n40000"
                   "\n50000"
                   "\n75000"
                   "\n100000"
                   "\n200000"
                   "\n400000"
                   "\n800000",
                   required=False, type=int)
    dp.add_argument('-l', '--link', help="link status, 0: down, 1: up",
                   required=False, type=int)
    dp.add_argument('-s', '--state', help="device state",
                   required=False, type=int)
    dp.add_argument('-f', '--features', help="Virtio Network Device features, can be:"
                   "\n1. +<bitmask|name>"
                   "\n   enable features, e.g +0x8000 or +VIRTIO_NET_F_MRG_RXBUF or +MRG_RXBUF"
                   "\n2. -<bitmask|name>"
                   "\n   disable features, "
                       "e.g -0x400000000 or -VIRTIO_NET_F_MRG_RXBUF or -MRG_RXBUF"
                   "\n3. <bitmask>"
                   "\n   set features, e.g 0x8100000300e7182f"
                   "\n4. <bitmask>?([+-]<bitmask|name>)*"
                   "\n   combined set+enable/disable features, "
                       "e.g 0x8100000300e7182f+0x400000000-MRG_RXBUF"
                   "\nsee supported features with 'virtnet list' command",
                   required=False, type=str)
    dp.add_argument('-o', '--supported_hash_types',
                   help="Supported hash type bitmask in hex: "
                   "\ne.g, 0x9 to enable both IPv4 and IPv6"
                   "\nVIRTIO_NET_HASH_TYPE_IPv4  1 << 0"
                   "\nVIRTIO_NET_HASH_TYPE_TCPv4 1 << 1"
                   "\nVIRTIO_NET_HASH_TYPE_UDPv4 1 << 2"
                   "\nVIRTIO_NET_HASH_TYPE_IPv6  1 << 3"
                   "\nVIRTIO_NET_HASH_TYPE_TCPv6 1 << 4"
                   "\nVIRTIO_NET_HASH_TYPE_UDPv6 1 << 5",
                   required=False, type=int_hex)
    dp.add_argument('-k', '--rss_max_key_size', help="RSS Maximum Key size",
                    required=False, type=int)
    dp.add_argument('-r', '--rx_mode', help="mode bits in hex, e.g. 0x1 "
                   "0:promisc 1:all-multi 2:all-uni 3:no-multi "
                   "4:no-uni 5:no-cast",
                   required=False, type=str)
    dp.add_argument('-n', '--msix_num', help="Modify VF number of MSIX", required=False, type=int)
    dp.add_argument('-q', '--max_queue_size', help="Modify max queue size of device", required=False, type=int)
    dp.add_argument('-d', '--dst_port', help="Modify ipv4 dst_port rules."
                    "\n-d {\"tcp\"|\"udp\"}:dst_port:rq_index"
                    "\n    packet matching tcp|udp dst_port is directed to specified rq."
                    "\n-d <\"tcp\"|\"udp\":0:xx>"
                    "\n    flush all tcp|udp rules on the device.",
                    required=False, type=str)
    dp.add_argument('-b', '--rx_dma_q_num', help="Modify max rx dma queue number",
                   required=False, type=int)
    dp.add_argument("-rxdim", '--rx_dim_config', help=argparse.SUPPRESS,
                   required=False, type=str, choices=['enable', 'disable'])
    dp.add_argument('-dc', '--drop_counter', help="enable/disable drop counter",
                   required=False, type=str, choices=['enable', 'disable'])
    dp.add_argument('-pkt_cnt', '--packet_counter', help="Enable/disable "
                    "packet counter stats via \"virtnet stats\"",
                   required=False, type=str, choices=['enable', 'disable'])
    dp.add_argument("-aarfs", '--aarfs_config', help="Enable/Disable auto Accelerated "
                    "Receive Flow Steering. ",
                    required=False, type=str, choices=['enable', 'disable'])
    dp.add_argument('-qp', '--max_queue_pairs', help="Modify maximum number of "
		    "queue pairs, each pair has 1 tx and 1 rx",
		    required=False, type=int)
    dp.add_argument("-dim", '--dim_config', help="Enable/Disable Dynamic "
                    "Interrupt Moderation(DIM) for both Rx and Tx VQs ",
                   required=False, type=str, choices=['enable', 'disable'])
    qp = sp.add_parser('queue', help='modify queue related options');
    qp.add_argument('-e', '--period_mode', help="period mode: event or cqe",
                   required=True, type=str, choices=['event', 'cqe'])
    qp.add_argument('-n', '--period', help="period (usec)",
                   required=True, type=int)
    qp.add_argument('-c', '--max_count', help="max count",
                   required=True, type=int)
    p.set_defaults(func=modify)
    dp.set_defaults(func=modify, subcmd='device')
    qp.set_defaults(func=modify, subcmd='queue')

    # Log
    p = subparsers.add_parser('log', help='set log level')
    p.add_argument('-l', '--level', help="log level: info/err/debug",
                    required=True, type=str, choices=['info', 'err', 'debug'])
    p.set_defaults(func=log_level)

    # Version
    p = subparsers.add_parser('version', help='show virtio net controller version info')
    p.set_defaults(func=version)

    # Restart
    p = subparsers.add_parser('restart', help='Do fast restart of controller without killing the service')
    p.set_defaults(func=restart)

    # Validate
    p = subparsers.add_parser('validate', help='validate configurations')
    p.add_argument('-f', '--file', help="validate virtnet config file, append path to virtnet.conf",
                   required=False, type=str)
    p.set_defaults(func=validate)

    # Update
    p = subparsers.add_parser('update', help='update controller')
    p.add_argument('-s', '--start', help="Start live update",
                   required=False, action="store_const", const=1, default=0)
    p.add_argument('-t', '--status', help="Check live update status",
                   required=False, action="store_const", const=1, default=0)
    p.set_defaults(func=update)

    # Health
    p = subparsers.add_parser('health', help='controller health utility')
    p.add_argument('-a', '--all', help='query health information for all virtnet devices, including '
                    'both PFs and VFs', required=False, action="store_const",
                    const=1, default=0)
    p.add_argument('-p', '--pf', help="pf ID", required=False, type=int)
    p.add_argument('-v', '--vf', help="vf ID", required=False, type=int)
    p.add_argument('-u', '--vuid', help="PF device VUID", required=False, type=str)
    sp = p.add_subparsers()
    hp = sp.add_parser('show', help='Show health info')

    p.set_defaults(func=health)
    hp.set_defaults(func=health, subcmd='show')

    # Debug
    p = subparsers.add_parser('debug', help='For debug purpose, cmds can be changed without notice')
    p.add_argument('-p', '--pf', help="pf ID", required=False, type=int)
    p.add_argument('-v', '--vf', help="vf ID", required=False, type=int)
    p.add_argument('-u', '--vuid', help="PF device VUID", required=False, type=str)
    sp = p.add_subparsers()
    dp = sp.add_parser('mem', help='Get memory info')
    dp.add_argument('-a', '--address', help="Memory address to get in format 0x10fe1b6000",
                    required=True, type=str)
    dp.add_argument('-n', '--num_dword', help="Number of dword to get",
                    required=True, type=int)
    qp = sp.add_parser('ring', help='Get ring info')
    qp.add_argument('-q', '--queue_id', help="Queue index of the device",
                    required=True, type=int)
    qp.add_argument('-r', '--ring_type', help="Ring type: avail/used/desc/desc_packed",
                    required=True, type=str, choices=['avail', 'used', 'desc', 'desc_packed'])
    qp.add_argument('-s', '--start', help="Start element of the ring",
                    required=False, type=int)
    qp.add_argument('-c', '--count', help="number of  element of the ring",
                    required=False, type=int)
    tp = sp.add_parser('trigger', help='Force trigger a special queue')
    tp.add_argument('-q', '--queue_id', help="Queue index of the device",
                    required=True, type=int)
    rp = sp.add_parser('options', help='Debug options to enable/disable')
    rp.add_argument("-dma_q_rec", '--dma_q_rec', help="Enable/Disable Dma queue recovery",
                    required=False, type=str, choices=['enable', 'disable'])
    rp.add_argument("-lm_log", '--lm_log', help="Enable/Disable Live migration buffer dump.",
                   required=False, type=str, choices=['enable', 'disable'])
    cp = sp.add_parser('coredump', help='Generate coredump if DPA crashed')
    ip = sp.add_parser('query', help='Query all virtnet devices debug info')
    mp = sp.add_parser('dev_mem', help='Dev memory dump')
    mp.add_argument('-a', '--address', help="Memory address to get in hex format with 0x prefix.",
                    required=True, type=str)
    mp.add_argument('-n', '--num_dword', help="Number of dword to get",
                    required=True, type=int)
    mp.add_argument('-f', '--file', help="Dump data into a file with absolute path.",
                    required=False, type=str)
    ap = sp.add_parser('aarfs', help='Accelerated Receive Flow Steering cmds')
    ap_hash = ap.add_subparsers()
    ap_hash_dump = ap_hash.add_parser('dump_hash', help='Dump aarfs hash info')
    ap_hash_dump.add_argument('-p', '--pf', help="pf ID", required=False, type=int)
    ap_mru_dump = ap_hash.add_parser('dump_mru', help='Dump aarfs mru info')
    ap_stats_dump = ap_hash.add_parser('dump_stats', help='Dump aarfs stats info')
    np = sp.add_parser('netdim', help='Get netdim stats')
    np.add_argument('-q', '--queue_id', help="Queue index of the device",
		    required=False, type=int)
    p.set_defaults(func=debug)
    dp.set_defaults(func=debug, subcmd='mem')
    qp.set_defaults(func=debug, subcmd='ring')
    tp.set_defaults(func=debug, subcmd='trigger')
    rp.set_defaults(func=debug, subcmd='options')
    cp.set_defaults(func=debug, subcmd='coredump')
    ip.set_defaults(func=debug, subcmd='query')
    mp.set_defaults(func=debug, subcmd='dev_mem')
    ap_hash_dump.set_defaults(func=debug, subcmd='dump_hash')
    ap_mru_dump.set_defaults(func=debug, subcmd='dump_mru')
    ap_stats_dump.set_defaults(func=debug, subcmd='dump_stats')
    np.set_defaults(func=debug, subcmd='netdim')

    #Stats
    p = subparsers.add_parser('stats', help='stats of virtnet device')
    p.add_argument('-p', '--pf', help="pf ID", required=False, type=int)
    p.add_argument('-v', '--vf', help="vf ID", required=False, type=int)
    p.add_argument('-u', '--vuid', help="PF device VUID", required=False,
		   type=str)
    p.add_argument('-q', '--queue_id', help="Queue index of the device",
                   required=False, type=int)
    p.set_defaults(func=stats)

    args = parser.parse_args()
    if args.called_rpc_name == "update" or args.called_rpc_name == "version" or \
       args.called_rpc_name == "restart":
        args.client = JsonRpcVirtnetClient(server_addr, manager_server_port, timeout)
    elif args.called_rpc_name == "health":
        args.client = JsonRpcVirtnetClient(server_addr, health_server_port, timeout)
    elif args.called_rpc_name == "validate":
        pass
    else:
        args.client = JsonRpcVirtnetClient(server_addr, server_port, timeout)
    if hasattr(args, 'func'):
        try:
            err = call_rpc_func(args)
            if err != None and err != 0:
                exit(1)
        except JsonRpcSnapException as ex:
            print(ex)
            exit(1)
    elif sys.stdin.isatty():
        # No arguments and no data piped through stdin
        parser.print_help()
        exit(1)
    else:
        execute_script(parser, args.client, sys.stdin)

if __name__ == "__main__":
    main()
