# This file is part of Scapy
# See http://www.secdev.org/projects/scapy for more information
# Copyright (C) Philippe Biondi <phil@secdev.org>
# This program is published under a GPLv2 license

from __future__ import print_function
import socket
from scapy.modules.six.moves.queue import Queue, Empty
from scapy.pipetool import Source, Drain, Sink
from scapy.config import conf
from scapy.compat import raw
from scapy.utils import PcapReader, PcapWriter
from scapy.automaton import recv_error
from scapy.consts import WINDOWS


class SniffSource(Source):
    """Read packets from an interface and send them to low exit.
     +-----------+
  >>-|           |->>
     |           |
   >-|  [iface]--|->
     +-----------+
"""

    def __init__(self, iface=None, filter=None, name=None):
        Source.__init__(self, name=name)
        self.iface = iface
        self.filter = filter

    def start(self):
        self.s = conf.L2listen(iface=self.iface, filter=self.filter)

    def stop(self):
        self.s.close()

    def fileno(self):
        return self.s.fileno()

    def check_recv(self):
        return True

    def deliver(self):
        try:
            self._send(self.s.recv())
        except recv_error:
            if not WINDOWS:
                raise


class RdpcapSource(Source):
    """Read packets from a PCAP file send them to low exit.
     +----------+
  >>-|          |->>
     |          |
   >-|  [pcap]--|->
     +----------+
"""

    def __init__(self, fname, name=None):
        Source.__init__(self, name=name)
        self.fname = fname
        self.f = PcapReader(self.fname)

    def start(self):
        print("start")
        self.f = PcapReader(self.fname)
        self.is_exhausted = False

    def stop(self):
        print("stop")
        self.f.close()

    def fileno(self):
        return self.f.fileno()

    def check_recv(self):
        return True

    def deliver(self):
        p = self.f.recv()
        print("deliver %r" % p)
        if p is None:
            self.is_exhausted = True
        else:
            self._send(p)


class InjectSink(Sink):
    """Packets received on low input are injected to an interface
     +-----------+
  >>-|           |->>
     |           |
   >-|--[iface]  |->
     +-----------+
"""

    def __init__(self, iface=None, name=None):
        Sink.__init__(self, name=name)
        if iface is None:
            iface = conf.iface
        self.iface = iface

    def start(self):
        self.s = conf.L2socket(iface=self.iface)

    def stop(self):
        self.s.close()

    def push(self, msg):
        self.s.send(msg)


class Inject3Sink(InjectSink):
    def start(self):
        self.s = conf.L3socket(iface=self.iface)


class WrpcapSink(Sink):
    """Packets received on low input are written to PCAP file
     +----------+
  >>-|          |->>
     |          |
   >-|--[pcap]  |->
     +----------+
"""

    def __init__(self, fname, name=None):
        Sink.__init__(self, name=name)
        self.f = PcapWriter(fname)

    def stop(self):
        self.f.flush()
        self.f.close()

    def push(self, msg):
        self.f.write(msg)


class UDPDrain(Drain):
    """UDP payloads received on high entry are sent over UDP
     +-------------+
  >>-|--[payload]--|->>
     |      X      |
   >-|----[UDP]----|->
     +-------------+
"""

    def __init__(self, ip="127.0.0.1", port=1234):
        Drain.__init__(self)
        self.ip = ip
        self.port = port

    def push(self, msg):
        from scapy.layers.inet import IP, UDP
        if IP in msg and msg[IP].proto == 17 and UDP in msg:
            payload = msg[UDP].payload
            self._high_send(raw(payload))

    def high_push(self, msg):
        from scapy.layers.inet import IP, UDP
        p = IP(dst=self.ip) / UDP(sport=1234, dport=self.port) / msg
        self._send(p)


class FDSourceSink(Source):
    """Use a file descriptor as source and sink
     +-------------+
  >>-|             |->>
     |             |
   >-|-[file desc]-|->
     +-------------+
"""

    def __init__(self, fd, name=None):
        Source.__init__(self, name=name)
        self.fd = fd

    def push(self, msg):
        self.fd.write(msg)

    def fileno(self):
        return self.fd.fileno()

    def deliver(self):
        self._send(self.fd.read())


class TCPConnectPipe(Source):
    """TCP connect to addr:port and use it as source and sink
     +-------------+
  >>-|             |->>
     |             |
   >-|-[addr:port]-|->
     +-------------+
"""
    __selectable_force_select__ = True

    def __init__(self, addr="", port=0, name=None):
        Source.__init__(self, name=name)
        self.addr = addr
        self.port = port
        self.fd = None

    def start(self):
        self.fd = socket.socket()
        self.fd.connect((self.addr, self.port))

    def stop(self):
        if self.fd:
            self.fd.close()

    def push(self, msg):
        self.fd.send(msg)

    def fileno(self):
        return self.fd.fileno()

    def deliver(self):
        try:
            msg = self.fd.recv(65536)
        except socket.error:
            self.stop()
            raise
        if msg:
            self._send(msg)


class TCPListenPipe(TCPConnectPipe):
    """TCP listen on [addr:]port and use first connection as source and sink ; send peer address to high output  # noqa: E501
     +------^------+
  >>-|    +-[peer]-|->>
     |   /         |
   >-|-[addr:port]-|->
     +-------------+
"""
    __selectable_force_select__ = True

    def __init__(self, addr="", port=0, name=None):
        TCPConnectPipe.__init__(self, addr, port, name)
        self.connected = False
        self.q = Queue()

    def start(self):
        self.connected = False
        self.fd = socket.socket()
        self.fd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.fd.bind((self.addr, self.port))
        self.fd.listen(1)

    def push(self, msg):
        if self.connected:
            self.fd.send(msg)
        else:
            self.q.put(msg)

    def deliver(self):
        if self.connected:
            try:
                msg = self.fd.recv(65536)
            except socket.error:
                self.stop()
                raise
            if msg:
                self._send(msg)
        else:
            fd, frm = self.fd.accept()
            self._high_send(frm)
            self.fd.close()
            self.fd = fd
            self.connected = True
            self._trigger(frm)
            while True:
                try:
                    self.fd.send(self.q.get(block=False))
                except Empty:
                    break


class TriggeredMessage(Drain):
    """Send a preloaded message when triggered and trigger in chain
     +------^------+
  >>-|      | /----|->>
     |      |/     |
   >-|-[ message ]-|->
     +------^------+
"""

    def __init__(self, msg, name=None):
        Drain.__init__(self, name=name)
        self.msg = msg

    def on_trigger(self, trigmsg):
        self._send(self.msg)
        self._high_send(self.msg)
        self._trigger(trigmsg)


class TriggerDrain(Drain):
    """Pass messages and trigger when a condition is met
     +------^------+
  >>-|-[condition]-|->>
     |      |      |
   >-|-[condition]-|->
     +-------------+
"""

    def __init__(self, f, name=None):
        Drain.__init__(self, name=name)
        self.f = f

    def push(self, msg):
        v = self.f(msg)
        if v:
            self._trigger(v)
        self._send(msg)

    def high_push(self, msg):
        v = self.f(msg)
        if v:
            self._trigger(v)
        self._high_send(msg)


class TriggeredValve(Drain):
    """Let messages alternatively pass or not, changing on trigger
     +------^------+
  >>-|-[pass/stop]-|->>
     |      |      |
   >-|-[pass/stop]-|->
     +------^------+
"""

    def __init__(self, start_state=True, name=None):
        Drain.__init__(self, name=name)
        self.opened = start_state

    def push(self, msg):
        if self.opened:
            self._send(msg)

    def high_push(self, msg):
        if self.opened:
            self._high_send(msg)

    def on_trigger(self, msg):
        self.opened ^= True
        self._trigger(msg)


class TriggeredQueueingValve(Drain):
    """Let messages alternatively pass or queued, changing on trigger
     +------^-------+
  >>-|-[pass/queue]-|->>
     |      |       |
   >-|-[pass/queue]-|->
     +------^-------+
"""

    def __init__(self, start_state=True, name=None):
        Drain.__init__(self, name=name)
        self.opened = start_state
        self.q = Queue()

    def start(self):
        self.q = Queue()

    def push(self, msg):
        if self.opened:
            self._send(msg)
        else:
            self.q.put((True, msg))

    def high_push(self, msg):
        if self.opened:
            self._send(msg)
        else:
            self.q.put((False, msg))

    def on_trigger(self, msg):
        self.opened ^= True
        self._trigger(msg)
        while True:
            try:
                low, msg = self.q.get(block=False)
            except Empty:
                break
            else:
                if low:
                    self._send(msg)
                else:
                    self._high_send(msg)


class TriggeredSwitch(Drain):
    r"""Let messages alternatively high or low, changing on trigger
     +------^------+
  >>-|-\    |    /-|->>
     |  [up/down]  |
   >-|-/    |    \-|->
     +------^------+
"""

    def __init__(self, start_state=True, name=None):
        Drain.__init__(self, name=name)
        self.low = start_state

    def push(self, msg):
        if self.low:
            self._send(msg)
        else:
            self._high_send(msg)
    high_push = push

    def on_trigger(self, msg):
        self.low ^= True
        self._trigger(msg)
