# encoding: utf-8
"""
netlink.py

Created by Thomas Mangin on 2015-03-31.
Copyright (c) 2009-2017 Exa Networks. All rights reserved.
License: 3-clause BSD. (See the COPYRIGHT file)
"""

import os
import socket
from struct import pack
from struct import unpack
from struct import calcsize
from collections import namedtuple

from exabgp.netlink import NetLinkError
from exabgp.netlink.sequence import Sequence
from exabgp.netlink.attributes import Attributes


try:
    getattr(socket, 'AF_NETLINK')
except AttributeError:
    raise ImportError('This module only works on unix version with netlink support')


# 0                   1                   2                   3
# 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
# |                                                               |
# |                   Netlink message header                      |
# |                                                               |
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
# |                                                               |
# |                  IP Service Template                          |
# |                                                               |
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
# |                                                               |
# |                  IP Service specific data in TLVs             |
# |                                                               |
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

# The Netlink message header

# 0                   1                   2                   3
# 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
# |                          Length                             |
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
# |            Type              |           Flags              |
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
# |                      Sequence Number                        |
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
# |                      Process ID (PID)                       |
# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+


class NetLink(object):
    _IGNORE_SEQ_FAULTS = True

    NETLINK_ROUTE = 0

    format = namedtuple('Message', 'format_type control_flags sequence pid data')
    netlink = socket.socket(socket.AF_NETLINK, socket.SOCK_RAW, NETLINK_ROUTE)

    class Header(object):
        # linux/netlink.h
        PACK = 'IHHII'
        LEN = calcsize(PACK)

    class Command(object):
        NLMSG_NOOP = 0x01
        NLMSG_ERROR = 0x02
        NLMSG_DONE = 0x03
        NLMSG_OVERRUN = 0x04

    class Flags(object):
        NLM_F_REQUEST = 0x01  # It is query message.
        NLM_F_MULTI = 0x02  # Multipart message, terminated by NLMSG_DONE
        NLM_F_ACK = 0x04  # Reply with ack, with zero or error code
        NLM_F_ECHO = 0x08  # Echo this query

        # Modifiers to GET query
        NLM_F_ROOT = 0x100  # specify tree root
        NLM_F_MATCH = 0x200  # return all matching
        NLM_F_DUMP = NLM_F_ROOT | NLM_F_MATCH
        NLM_F_ATOMIC = 0x400  # atomic GET

        # Modifiers to NEW query
        NLM_F_REPLACE = 0x100  # Override existing
        NLM_F_EXCL = 0x200  # Do not touch, if it exists
        NLM_F_CREATE = 0x400  # Create, if it does not exist
        NLM_F_APPEND = 0x800  # Add to end of list

    errors = {
        Command.NLMSG_ERROR: 'netlink error',
        Command.NLMSG_OVERRUN: 'netlink overrun',
    }

    @classmethod
    def encode(cls, format_type, sequence, control_flags, pid, body, attributes):
        attrs = Attributes.encode(attributes)
        length = cls.Header.LEN + len(attrs) + len(body)
        return pack(cls.Header.PACK, length, format_type, control_flags, sequence, pid) + body + attrs

    @classmethod
    def decode(cls, data):
        while data:
            length, format_type, control_flags, sequence, pid = unpack(cls.Header.PACK, data[: cls.Header.LEN])
            if len(data) < length:
                raise NetLinkError("Buffer underrun")
            yield cls.format(format_type, control_flags, sequence, pid, data[cls.Header.LEN : length])
            data = data[length:]

    # pack('Bxxx', family),
    # family=socket.AF_UNSPEC,

    @classmethod
    def send(cls, format_type, control_flags, family=socket.AF_UNSPEC, attributes=None):
        sequence = Sequence()
        pid = os.getpid()

        if attributes is None:
            attributes = {}

        message = cls.encode(format_type, sequence, control_flags, pid, pack('Bxxx', family), attributes)

        cls.netlink.send(message)

        while True:
            response = cls.netlink.recv(640000)
            for response_type, response_flags, response_sequence, response_pid, response_data in cls.decode(response):
                if response_sequence != sequence:
                    if cls._IGNORE_SEQ_FAULTS:
                        continue
                    raise NetLinkError("netlink sequence mismatch")
                elif response_pid != pid:
                    raise NetLinkError("netlink pid mismatch")
                # elif response_flags != control_flags:
                # 	raise NetLinkError("netlink flags mismatch")
                elif response_type in cls.errors:
                    raise NetLinkError(cls.errors[response_type], message, response)

                if response_type == NetLink.Command.NLMSG_DONE:
                    raise StopIteration()
                yield response_data