#!/usr/bin/env python
"""Change tx power parameter.

Derived from a Python test of libnl available here:
    https://github.com/Robpol86/libnl/blob/7186e04/tests/test_nl80211.py#L29

This script is a showcase for the Linux Netlink libnl Python library, ported
from the C library with the same name. Using libnl, Python scripts/applications
can communicate with the Linux kernel and device drivers directly just like
C/C++ programs do, without having to call a binary on the system.

More information about the Python libnl is available here:
    https://github.com/Robpol86/libnl

Debug messages are available with the -v option, and even more debug messages
are available by setting the NLCB environment variable to either 'verbose' or
'debug' like so:
    NLCB=verbose change_wifi_tx_power.py -v
    NLCB=debug change_wifi_tx_power.py 10 wlan0

Usage:
    change_wifi_tx_power.py [options] [<interface>] [<tx-power>]
    change_wifi_tx_power.py -h | --help

Options:
    -v --verbose    Print debug messages to stderr.
"""

from __future__ import print_function

import fcntl
import logging
import signal
import socket
import struct
import sys

from docopt import docopt
#from terminaltables import AsciiTable

from libnl.attr import nla_data, nla_get_string, nla_get_u32, nla_parse, nla_put_u32, nla_get_u64
from libnl.error import errmsg
from libnl.genl.ctrl import genl_ctrl_resolve
from libnl.genl.genl import genl_connect, genlmsg_attrdata, genlmsg_attrlen, genlmsg_put
from libnl.handlers import NL_CB_CUSTOM, NL_CB_VALID, NL_SKIP
from libnl.linux_private.genetlink import genlmsghdr
from libnl.linux_private.netlink import NLM_F_DUMP
from libnl.msg import nlmsg_alloc, nlmsg_data, nlmsg_hdr
from libnl.nl import nl_recvmsgs_default, nl_send_auto
from libnl.nl80211 import nl80211
from libnl.socket_ import nl_socket_alloc, nl_socket_modify_cb

OPTIONS = docopt(__doc__) if __name__ == '__main__' else dict()


def error(message, code=1):
    """Print an error message to stderr and exits with a status of 1 by default."""
    if message:
        print('ERROR: {0}'.format(message), file=sys.stderr)
    else:
        print(file=sys.stderr)
    sys.exit(code)


def main(tx_power):
    """Main function called upon script execution."""
    # First get the wireless interface index.
    if OPTIONS['<interface>']:
        pack = struct.pack('16sI', OPTIONS['<interface>'].encode('ascii'), 0)
        sk = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            info = struct.unpack('16sI', fcntl.ioctl(sk.fileno(), 0x8933, pack))
        except OSError:
            return error('Wireless interface {0} does not exist.'.format(OPTIONS['<interface>']))
        finally:
            sk.close()
        if_index = int(info[1])
    else:
        if_index = -1

    # Then open a socket to the kernel. Same one used for sending and receiving.
    sk = nl_socket_alloc()  # Creates an `nl_sock` instance.
    ret = genl_connect(sk)  # Create file descriptor and bind socket.
    if ret < 0:
        reason = errmsg[abs(ret)]
        return error('genl_connect() returned {0} ({1})'.format(ret, reason))

    # Now get the nl80211 driver ID. Handle errors here.
    driver_id = genl_ctrl_resolve(sk, b'nl80211')  # Find the nl80211 driver ID.
    if driver_id < 0:
        reason = errmsg[abs(driver_id)]
        return error('genl_ctrl_resolve() returned {0} ({1})'.format(driver_id, reason))

    # Setup the Generic Netlink message.
    msg = nlmsg_alloc()  # Allocate a message.
    if OPTIONS['<interface>']:
        genlmsg_put(msg, 0, 0, driver_id, 0, 0, nl80211.NL80211_CMD_SET_WIPHY, 0)  # Tell kernel: send iface info.
        nla_put_u32(msg, nl80211.NL80211_ATTR_IFINDEX, if_index)  # This is the interface we care about.
        nla_put_u32(msg, nl80211.NL80211_ATTR_WIPHY_TX_POWER_SETTING, 2)  # 
        nla_put_u32(msg, nl80211.NL80211_ATTR_WIPHY_TX_POWER_LEVEL, int(tx_power) * 100)  # This is the interface we care about.
    else:
        # Interface name should be provided 
        print("Interface name should be provided")

    # Now send the message to the kernel, and get its response, automatically calling the callback.
    ret = nl_send_auto(sk, msg)
    if ret < 0:
        reason = errmsg[abs(ret)]
        return error('nl_send_auto() returned {0} ({1})'.format(ret, reason))
    print('Sent {0} bytes to the kernel.'.format(ret))
    ret = nl_recvmsgs_default(sk)  # Blocks until the kernel replies. Usually it's instant.
    if ret < 0:
        reason = errmsg[abs(ret)]
        return error('nl_recvmsgs_default() returned {0} ({1})'.format(ret, reason))


def setup_logging():
    """Called when __name__ == '__main__' below. Sets up logging library.

    All logging messages go to stderr, from DEBUG to CRITICAL. This script uses print() for regular messages.
    """
    fmt = 'DBG<0>%(pathname)s:%(lineno)d  %(funcName)s: %(message)s'

    handler_stderr = logging.StreamHandler(sys.stderr)
    handler_stderr.setFormatter(logging.Formatter(fmt))

    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)
    root_logger.addHandler(handler_stderr)


if __name__ == '__main__':
    signal.signal(signal.SIGINT, lambda *_: sys.exit(0))  # Properly handle Control+C
    if OPTIONS.get('--verbose'):
        setup_logging()
    if OPTIONS['<tx-power>']:
        print("Tx Power",OPTIONS['<tx-power>'])
        main(OPTIONS['<tx-power>'])

