# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import dpkt
import socket
import struct

from lansim import tools


# An initial set of Protocol to Hardware address mappings.
_ARP_INITIAL_CACHE = {
    # Broadcast address:
    socket.inet_aton('255.255.255.255'): tools.inet_hwton('FF:FF:FF:FF:FF:FF'),
}


class SimpleHostError(Exception):
    """A SimpleHost generic error."""


class SimpleHost(object):
    """A simple host supporting IPv4.

    This class is useful as a base clase to implement other hosts. It supports
    a single IPv4 address.
    """
    def __init__(self, sim, hw_addr, ip_addr):
        """Creates the host and associates it with the given NetworkBridge.

        @param sim: The Simulator interface where this host lives.
        @param hw_addr: Hex or binary representation of the Ethernet address.
        @param ip_addr: The IPv4 address. For example: "10.0.0.1".
        """
        self._sim = sim
        self._hw_addr = hw_addr
        self._ip_addr = ip_addr
        self._bin_hw_addr = tools.inet_hwton(hw_addr)
        self._bin_ip_addr = socket.inet_aton(ip_addr)
        # arp cache: Protocol to Hardware address resolution cache.
        self._arp_cache = dict(_ARP_INITIAL_CACHE)
        # Reply to broadcast ARP requests.
        rule = {
            "dst": "\xff" * 6, # Broadcast HW addr.
            "arp.pln": 4, # Protocol Addres Length is 4 (IP v4).
            "arp.op": dpkt.arp.ARP_OP_REQUEST,
            "arp.tpa": self._bin_ip_addr}
        sim.add_match(rule, self.arp_request)

        # Reply to unicast ARP requests.
        rule["dst"] = self._bin_hw_addr
        sim.add_match(rule, self.arp_request)

        # Mappings used for TCP traffic forwarding.
        self._tcp_fwd_in = {}
        self._tcp_fwd_out = {}
        self._tcp_fwd_ports = {}

    @property
    def ip_addr(self):
        """Returns the host IPv4 address."""
        return self._ip_addr


    @property
    def simulator(self):
        """Returns the Simulator instance where this host runs on."""
        return self._sim


    def arp_request(self, pkt):
        """Sends the ARP_REPLY matching the request.

        @param pkt: a dpkt.Packet with the ARP_REQUEST.
        """
        # Update the local ARP cache whenever we get a request.
        self.add_arp(hw_addr=pkt.arp.sha, ip_addr=pkt.arp.spa)

        arp_resp = dpkt.arp.ARP(
            op = dpkt.arp.ARP_OP_REPLY,
            pln = 4,
            tpa = pkt.arp.spa, # Target Protocol Address.
            tha = pkt.arp.sha, # Target Hardware Address.
            spa = self._bin_ip_addr, # Source Protocol Address.
            sha = self._bin_hw_addr) # Source Hardware Address.
        eth_resp = dpkt.ethernet.Ethernet(
            dst = pkt.arp.sha,
            src = self._bin_hw_addr,
            type = dpkt.ethernet.ETH_TYPE_ARP,
            data = arp_resp)
        self._sim.write(eth_resp)


    def add_arp(self, hw_addr, ip_addr):
        """Maps the ip_addr to a given hw_addr.

        This is useful to send IP packets with send_ip() to hosts that haven't
        comunicate with us yet.

        @param hw_addr: The network encoded corresponding Ethernet address.
        @param ip_addr: The network encoded IPv4 address.
        """
        self._arp_cache[ip_addr] = hw_addr


    def _resolve_mac_address(self, ip_addr):
        """Resolves the hw_addr of an IP address locally when it is known.

        This method uses the information gathered from received ARP requests and
        locally added mappings with add_arp(). It also knows how to resolve
        multicast addresses.

        @param ip_addr: The IP address to resolve encoded in network format.
        @return: The Hardware address encoded in network format or None
        if unknown.
        @raise SimpleHostError if the MAC address for ip_addr is unknown.
        """
        # From RFC 1112 6.4:
        #  An IP host group address is mapped to an Ethernet multicast address
        #  by placing the low-order 23-bits of the IP address into the low-order
        #  23 bits of the Ethernet multicast address 01-00-5E-00-00-00 (hex).
        #  Because there are 28 significant bits in an IP host group address,
        #  more than one host group address may map to the same Ethernet
        #  multicast address.
        int_ip_addr, = struct.unpack('!I', ip_addr)
        if int_ip_addr & 0xF0000000 == 0xE0000000: # Multicast IP address
            int_hw_ending = int_ip_addr & ((1 << 23) - 1) | 0x5E000000
            return '\x01\x00' + struct.pack('!I', int_hw_ending)
        if ip_addr in self._arp_cache:
            return self._arp_cache[ip_addr]
        # No address found.
        raise SimpleHostError("Unknown destination IP host.")


    def send_ip(self, pkt):
        """Sends an IP packet.

        The source IP address and the hardware layer is automatically filled.
        @param pkt: A dpkg.ip.IP packet.
        @raise SimpleHostError if the MAC address for ip_addr is unknown.
        """
        hw_dst = self._resolve_mac_address(pkt.dst)

        pkt.src = self._bin_ip_addr
        # Set the packet length and force to recompute the checksum.
        pkt.len = len(pkt)
        pkt.sum = 0
        hw_pkt = dpkt.ethernet.Ethernet(
            dst = hw_dst,
            src = self._bin_hw_addr,
            type = dpkt.ethernet.ETH_TYPE_IP,
            data = pkt)
        return self._sim.write(hw_pkt)


    def tcp_forward(self, port, dest_addr, dest_port):
        """Forwards all the TCP/IP traffic on a given port to another host.

        This method makes all the incoming traffic for this host on a particular
        port be redirected to dest_addr:dest_port. This allows us to use the
        kernel's network stack to handle that traffic.

        @param port: The TCP port on this simulated host.
        @param dest_addr: A host IP address on the same network in plain text.
        @param dest_port: The TCP port on the destination host.
        """
        if not self._tcp_fwd_ports:
            # Lazy initialization.
            self._sim.add_match({
                'ip.dst': self._bin_ip_addr,
                'ip.p': dpkt.ip.IP_PROTO_TCP}, self._handle_tcp_forward)

        self._tcp_fwd_ports[port] = socket.inet_aton(dest_addr), dest_port


    def _tcp_pick_port(self, dhost, dport):
        """Picks a new unused source TCP port on the host."""
        for p in range(1024, 65536):
            if (dhost, dport, p) in self._tcp_fwd_out:
                continue
            if p in self._tcp_fwd_ports:
                continue
            return p
        raise SimpleHostError("Too many connections.")


    def _handle_tcp_forward(self, pkt):
        # Source from:
        shost = pkt.ip.src
        sport = pkt.ip.tcp.sport
        dport = pkt.ip.tcp.dport

        ### Handle responses from forwarded traffic back to the sender (out).
        if (shost, sport, dport) in self._tcp_fwd_out:
            fhost, fport, oport = self._tcp_fwd_out[(shost, sport, dport)]
            # Redirect the packet
            pkt.ip.tcp.sport = oport
            pkt.ip.tcp.dport = fport
            pkt.ip.dst = fhost
            pkt.ip.tcp.sum = 0 # Force checksum
            self.send_ip(pkt.ip)
            return

        ### Handle incoming traffic to a local forwarded port (in).
        if dport in self._tcp_fwd_ports:
            # Forward to:
            fhost, fport = self._tcp_fwd_ports[dport]

            ### Check if it is an existing connection.
            # lport: The port from where we send data out.
            if (shost, sport, dport) in self._tcp_fwd_in:
                lport = self._tcp_fwd_in[(shost, sport, dport)]
            else:
                # Pick a new local port on our side.
                lport = self._tcp_pick_port(fhost, fport)
                self._tcp_fwd_in[(shost, sport, dport)] = lport
                self._tcp_fwd_out[(fhost, fport, lport)] = (shost, sport, dport)

            # Redirect the packet
            pkt.ip.tcp.sport = lport
            pkt.ip.tcp.dport = fport
            pkt.ip.dst = fhost
            pkt.ip.tcp.sum = 0 # Force checksum
            self.send_ip(pkt.ip)
            return


    def socket(self, family, sock_type):
        """Creates an asynchronous socket on the simulated host.

        This method creates an asynchronous socket object that can be used to
        receive and send packets. This module only supports UDP sockets.

        @param family: The socket family, only AF_INET is supported.
        @param sock_type: The socket type, only SOCK_DGRAM is supported.
        @return: an UDPSocket object. See UDPSocket documentation for details.
        @raise SimpleHostError if socket family and type is not supported.
        """
        if family != socket.AF_INET:
            raise SimpleHostError("socket family not supported.")
        if sock_type != socket.SOCK_DGRAM:
            raise SimpleHostError("socket type not supported.")

        return UDPSocket(self)


class UDPSocket(object):
    """An asynchronous UDP socket interface.

    This UDP socket interface provides a way to send and received UDP messages
    on an asynchronous way. This means that the socket doesn't have a recv()
    method as a normal socket would have, since the simulation is event driven
    and a callback should not block its execution. See the listen() method to
    see how to receive messages from this socket.

    This interface is used by modules outside lansim to interact with lansim
    in a way that can be ported to other different backends. For example, this
    same interface could be implemented using the Python's socket module and
    the real kernel stack.
    """
    def __init__(self, host):
        """Initializes the UDP socket.

        To be used for receiving packets, listen() must be called.

        @param host: A SimpleHost object.
        """
        self._host = host
        self._sim = host.simulator
        self._port = None


    def __del__(self):
        self.close()


    def listen(self, ip_addr, port, recv_callback):
        """Bind and listen on the ip_addr:port.

        Calls recv_callback(pkt, src_addr, src_port) every time an UDP frame
        is received. src_addr and src_port are passed with the source IPv4
        (as in '192.168.0.2') and the sender port number.

        This function can only be called once, since the socket can't be
        reassigned.

        @param ip_addr: Local destination ip_addr. If None, the Host's IPv4
        address is used, for example '224.0.0.251' or '192.168.0.1'.
        @param port: Local destination port number.
        @param recv_callback: A callback function that accepts three
        arguments, the received string, the sender IPv4 address and the
        sender port number.
        """
        if ip_addr is None:
            ip_addr = self._host.ip_addr()
        self._port = port

        # Binds all the traffic to the provided callback converting the
        # single argument callback to the multiple argument.
        self._sim.add_match({
            "ip.dst": socket.inet_aton(ip_addr),
            "ip.udp.dport": port},
            lambda pkt: recv_callback(pkt.ip.udp.data,
                                      socket.inet_ntoa(pkt.ip.src),
                                      pkt.ip.udp.sport))


    def send(self, data, ip_addr, port):
        """Send an UDP message with the data string to ip_addr:port.

        @param data: Any string small enough to fit in a single UDP packet.
        @param ip_addr: Destination IPv4 address.
        @param port: Destination UDP port number.
        """
        pkt_udp = dpkt.udp.UDP(
            dport = port,
            sport = self._port if self._port != None else 0,
            data = data
        )
        # dpkt doesn't set the Length field on UDP packets according to RFC 768.
        pkt_udp.ulen = len(pkt_udp.pack_hdr()) + len(str(pkt_udp.data))

        pkt_ip = dpkt.ip.IP(
            dst = socket.inet_aton(ip_addr),
            ttl = 255, # The comon value for IP packets.
            off = dpkt.ip.IP_DF, # Don't frag.
            p = dpkt.ip.IP_PROTO_UDP,
            data = pkt_udp
        )
        self._host.send_ip(pkt_ip)


    def close(self):
        """Closes the socket and disconnects the bound callback."""
        #TODO(deymo): Remove the add_match rule added on listen().