#!/usr/bin/env python3
"""Python netcat implementation."""

# TODO: change docstring to this format: https://sphinxcontrib-napoleon.readthedocs.io/en/latest/
#
# Behaviour Server
# 1. (tty and non-tty mode) TCP Server should always quit if client disconnects
# 2. (tty and non-tty mode) TCP Server should always quit if connection is gone/broken
# 3. (tty and non-tty mode) UDP Will always stay open (can't determine if connection is gone)

from __future__ import print_function
from abc import ABCMeta, abstractmethod
from subprocess import Popen, PIPE
import argparse
import atexit
import logging
import os
import re
import socket
import subprocess
import sys
import threading
import time

# Python 2 + Python 3 support
ABC = ABCMeta("ABC", (object,), {"__slots__": ()})

# TODO: Find windows import
if os.name != "nt":
    import select

# -------------------------------------------------------------------------------------------------
# GLOBALS
# -------------------------------------------------------------------------------------------------

APPNAME = "pwncat"
APPREPO = "https://github.com/cytopia/pwncat"
VERSION = "0.0.8-alpha"

# Custom loglevel numer for TRACE
LOGLEVEL_TRACE_NUM = 9

# Default timeout for timeout-based sys.stdin and socket.recv
TIMEOUT_READ_STDIN = 0.1
TIMEOUT_RECV_SOCKET = 0.1
TIMEOUT_RECV_SOCKET_RETRY = 2


# #################################################################################################
# #################################################################################################
# ###
# ###   1 / 3   L I B R A R Y   C L A S S E S
# ###
# #################################################################################################
# #################################################################################################

# -------------------------------------------------------------------------------------------------
# CLASS: StringEncoder
# -------------------------------------------------------------------------------------------------
class StringEncoder(object):
    """
    Takes care about Python 2/3 string encoding/decoding.

    This allows to parse all string/byte values internally between all
    classes or functions as strings to keep full Python 2/3 compat.
    """

    # https://stackoverflow.com/questions/606191/27527728#27527728
    codec = "cp437"

    def __init__(self):
        self.py3 = sys.version_info >= (3, 0)

    def encode(self, data):
        """Convert string into a byte type for Python3."""
        if self.py3:
            data = data.encode(self.codec)
        return data

    def decode(self, data):
        """Convert bytes into a string type for Python3."""
        if self.py3:
            data = data.decode(self.codec)
        return data


# -------------------------------------------------------------------------------------------------
# ABSTRCT CLASS: AbstractSocket
# -------------------------------------------------------------------------------------------------
class AbstractSocket(object):
    """Abstract class which provides TCP, UDP and IPv4, IPv6 Socket funcionality."""

    sock = None  # server binding socket (until accept())
    conn = None  # client/server communication socket

    # The instance role must be 'server' or 'client' and
    # is used to determine how to reconnect broken connections.
    # Either listen again (tcp-only) or re-connect to upstream.
    role = None  # Must be 'server' or 'client'

    # Specify a float of seconds for the socket timeout.
    # A default of None means it is blocking.
    recv_timeout = None

    # If no data is received, the socket will timeout and counts subsequent timeouts
    # with no data received. If 'recv_timeout_retry' many timeouts occured, the socket
    # will stop reading data.
    # The internal counter is reset as soon as data is received again.
    recv_timeout_retry = 0

    options = {
        "bufsize": 1024,  # Receive buffer size
        "backlog": 0,  # Listen backlog
        "nodns": False,  # Do not resolve hostname
        "udp": False,  # Is TCP or UDP?
        "http": False,  # Use HTTP instead
        "https": False,  # Use HTTPS instead
        "keep_open": False,  # Keep server open for new connections
        "rebind": False,  # False (never), True (indefinite) or int for how many times to rebind
        "reconn": False,  # False (never), True (indefinite) or int for how many times to reconnect
        "rebind_wait": 0,  # Time in seconds between re-binds
        "reconn_wait": 0,  # Time in seconds between reconnects
        "rebind_robin": [],  # Ports to round-robind for bindings
        "reconn_robin": [],  # Ports to round-robin for reconnects
        "wait": False,  # Close client if it is idle for x many seconds
        "udp_ping_init": False,  # Send initial UDP ping packet
        "udp_ping_intvl": False,  # Interval in sec for UDP client to ping server
        "udp_ping_word": "\0",  # The char/string to send as UDP ping probe
        "udp_ping_robing": [],  # Ports to round-robing during UDP ping probes
        "safe_word": False,  # Once this is received, the application quits
    }

    # Store the address of the remote end.
    # If we are in server role and running in UDP mode,
    # it must wait for the client to connect first in order
    # to retrieve its addr and port in order to be able
    # to send data back to it.
    remote_addr = None
    remote_port = None

    # ------------------------------------------------------------------------------
    # Constructor / Destructor
    # ------------------------------------------------------------------------------
    def __init__(self, encoder, role, recv_timeout, recv_timeout_retry, options={}):
        """Constructor."""
        assert type(self) is not AbstractSocket, "AbstractSocket cannot be instantiated directly."
        assert role in ["server", "client"], "The role must be 'server' or 'client'."

        self.log = logging.getLogger(__name__)
        self.enc = encoder
        self.role = role

        self.recv_timeout = recv_timeout
        self.recv_timeout_retry = recv_timeout_retry

        # Ensure to provide all items (be explicit - currently helps to find bugs)
        for index in self.options:
            assert index in options, "Provided options are missing key: {}".format(index)
        # Ensure to only provide option itenms that are defined/exist (helps to find bugs)
        for index in options:
            assert index in self.options, "Provided an invalid option item"
            self.options[index] = options[index]

        # Register destructor
        atexit.register(self.__exit__)

    def __exit__(self):
        """Destructor."""

        self.log.trace("Closing 'sock' socket")
        self.__close_socket(self.sock)
        self.sock = None
        self.log.trace("Closing 'conn' socket")
        self.__close_socket(self.conn)
        self.conn = None

    # ------------------------------------------------------------------------------
    # Private Functions
    # ------------------------------------------------------------------------------
    def __close_socket(self, sock):
        """Close a socket."""
        try:
            # (SHUT_RD)   0 = Done receiving (disallows receiving)
            # (SHUT_WR)   1 = Done sending (disallows sending)
            # (SHUT_RDWR) 2 = Both
            sock.shutdown(socket.SHUT_RDWR)
        except (AttributeError, OSError, socket.error):
            self.log.trace("Could not shutdown socket")
            pass
        try:
            sock.close()
        except (AttributeError, OSError, socket.error):
            self.log.trace("Could not shutdown socket")

    def __reconnect(self):
        """Reconnect to a server if upstream has gone."""
        self.__close_socket(self.conn)
        self.__close_socket(self.sock)
        self.create_socket()
        self.conn = self.sock
        if not self.connect():
            self.__reconnect_to_server()

    def __reaccept_from_client(self):
        """Ensure the server is able to keep connection open by re-accepting new clients."""
        # Only for server
        assert self.role == "server", "Only the role 'server' can accept connections."
        # Do not re-accept for UDP
        assert not self.options["udp"], "This should have been caught during arg check."

        # [NO] Do not re-accept
        if not self.options["keep_open"]:
            self.log.info("No automatic re-accept specified. Shutting down.")
            return False
        # [YES] Re-accept indefinitely
        self.log.info("Re-accepting new clients")
        self.__close_socket(self.conn)
        self.accept()
        return True

    def __reconnect_to_server(self):
        """Ensure the client re-connects to the remote server, if the remote server hang up."""
        # Only for Clients
        assert self.role == "client", "Only the role 'client' can re-connect."
        # Do not re-connect with UDP
        assert not self.options["udp"], "This should have been caught during arg check."

        # [NO] Never re-connect
        if type(self.options["reconn"]) is bool and not self.options["reconn"]:
            self.log.info("No automatic reconnect specified. Shutting down.")
            return False
        # [YES] Always re-connect indefinitely
        if type(self.options["reconn"]) is bool and self.options["reconn"]:
            self.log.info(
                "Reconnecting in {} sec (indefinitely)".format(self.options["reconn_wait"])
            )
            time.sleep(self.options["reconn_wait"])
            self.__reconnect()
            return True
        # [YES] Re-connect x many times
        if self.options["reconn"] > 0:
            self.log.info(
                "Reconnecting in {} sec ({} more times left)".format(
                    self.options["reconn_wait"], self.options["reconn"]
                )
            )
            self.options["reconn"] -= 1
            time.sleep(self.options["reconn_wait"])
            self.__reconnect()
            return True
        # [NO] Re-connect count is used up
        self.log.info("Reconnect count is used up. Shutting down.")
        return False

    # ------------------------------------------------------------------------------
    # Helper Functions
    # ------------------------------------------------------------------------------
    def gethostbyname(self, host, port, family):
        """Translate hostname into IP address."""
        socktype = socket.SOCK_DGRAM if self.options["udp"] else socket.SOCK_STREAM
        flags = 0
        if self.options["nodns"]:
            flags = socket.AI_NUMERICHOST
        try:
            self.log.debug("Resolving hostname: {}".format(host))
            (family, socktype, proto, canonname, sockaddr) = socket.getaddrinfo(
                host, port, socket.AF_INET, socktype, flags
            )[0]
        except socket.gaierror as error:
            self.log.error("Resolve Error: {}".format(error))
            sys.exit(1)
        self.log.debug("Resolved hostname:  {}".format(sockaddr[0]))
        return sockaddr[0]

    def create_socket(self):
        """Create TCP or UDP socket."""
        try:
            if self.options["udp"]:
                self.log.debug("Creating UDP socket")
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            else:
                self.log.debug("Creating TCP socket")
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        except socket.error as error:
            self.log.error("Failed to create the socket: {}".format(error))
            sys.exit(1)
        # Get around the "[Errno 98] Address already in use" error, if the socket is still in wait
        # we instruct it to reuse the address anyway.
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    def bind(self, addr, port):
        """Bind the socket to an address."""
        try:
            self.log.debug("Binding socket to {}:{}".format(addr, port))
            self.sock.bind((addr, port))
        except (OverflowError, OSError, socket.error) as error:
            self.log.error("Binding socket to {}:{} failed: {}".format(addr, port, error))
            sys.exit(1)

    def listen(self):
        """Listen for connections made to the socket."""
        try:
            self.log.debug("Listening with backlog={}".format(self.options["backlog"]))
            self.sock.listen(self.options["backlog"])
        except socket.error as error:
            self.log.error("Listening failed: {}".format(error))
            sys.exit(1)

    def accept(self):
        """Accept a connection."""
        try:
            self.log.debug("Waiting for TCP client")
            self.conn, client = self.sock.accept()
            addr, port = client
            self.remote_addr = addr
            self.remote_port = port
            self.log.info("Client connected from {}:{}".format(addr, port))
        except (socket.gaierror, socket.error) as error:
            self.log.error("Accept failed: {}".format(error))
            sys.exit(1)

    def connect(self):
        """Connect to a remote socket at given address and port (TCP-only)."""
        try:
            self.log.debug("Connecting to {}:{}".format(self.remote_addr, self.remote_port))
            self.sock.connect((self.remote_addr, self.remote_port))
            return True
        except socket.error as error:
            self.log.error(
                "Connecting to {}:{} failed: {}".format(self.remote_addr, self.remote_port, error)
            )
            return False

    # ------------------------------------------------------------------------------
    # Send / Receive Functions
    # ------------------------------------------------------------------------------
    def send(self, data):
        """Send data."""
        # In case of sending data back to an udp client we need to wait
        # until the client has first connected and told us its addr/port
        if self.options["udp"] and self.remote_addr is None and self.remote_port is None:
            self.log.warning("UDP client has not yet connected. Queueing message")
            while self.remote_addr is None and self.remote_port is None:
                time.sleep(0.1)  # Less wastefull than using 'pass'

        curr = 0  # bytes send during one loop iteration
        send = 0  # total bytes send
        size = len(data)  # bytes of data that needs to be send
        data = self.enc.encode(data)
        assert size == len(data), "Encoding messed up string length, might need to do len() after."

        # Loop until all bytes have been send
        while send < size:
            try:
                self.log.debug(
                    "Trying to send {} bytes to {}:{}".format(
                        size - send, self.remote_addr, self.remote_port
                    )
                )
                self.log.trace("Trying to send: {}".format(repr(data)))
                if self.options["udp"]:
                    curr = self.conn.sendto(data, (self.remote_addr, self.remote_port))
                    send += curr
                else:
                    curr = self.conn.send(data)
                    send += curr
                if curr == 0:
                    self.log.error("No bytes send during loop round.")
                    return
                # Remove 'curr' many bytes from data for the next round
                data = data[curr:]
                self.log.debug(
                    "Sent {} bytes to {}:{} ({} bytes remaining)".format(
                        curr, self.remote_addr, self.remote_port, size - send
                    )
                )
            except socket.error as error:
                if error.errno == socket.errno.EPIPE:
                    self.log.error("TODO:Add desc. Socket error({}): {}".format(error.errno, error))
                    return
                # Most likely nothing to see here??
                # FIXME: TODO: Need to re-accepd new client
                self.log.error("TODO:Add desc. Socket Error: {}".format(error))
                if self.role == "server":
                    if self.__reaccept_from_client():
                        continue
                self.log.warning("Shutdown")
                return
            except (OSError) as error:
                self.log.error("Socket OS Error: {}".format(error))
                return

    def receive(self, ssig):
        """
        Generator function to receive data endlessly by yielding it.

        :param function interrupter: A Func that returns True/False to tell us to stop or not.
        """
        # Set current receive timeout
        self.conn.settimeout(self.recv_timeout)
        self.log.trace("Socket Timeout: {}".format(self.recv_timeout))
        # Counts how many times we had a ready timeout for later to decide
        # if we exceeded maximum retires
        curr_recv_timeout_retry = 0

        while True:
            # Ensure to signal that we do not stop receiving data
            # if ssig.has_stop():
            #    self.log.debug("Interrupt has been requested for receive()")
            #    return

            if self.conn is None:
                self.log.error("Exit. Socket is gone in receive()")
                ssig.raise_stop()
                return

            # Non-blocking socket with timeout. If the timeout threshold is hit,
            # it will throw an socket.timeout exception. This is required to see if other
            # threads have been terminated already.
            try:
                # https://manpages.debian.org/buster/manpages-dev/recv.2.en.html
                (byte, addr) = self.conn.recvfrom(self.options["bufsize"])

            # [1/5] Finished receiving all data
            # NOTE: This is the place where we can do any checks in between reads as the
            # socket has been changed from blocking to time-out based.
            # NOTE: This is also the place, where we quit in case --wait was specified.
            except socket.timeout:
                # Let's ask the interrupter() function if we should terminate?
                if not ssig.has_stop():
                    # No action required, continue the loop and read again.
                    continue
                self.log.debug("Interrupt has been requested for receive()")
                # Other threads are done. Let's try to read a few more times before
                # returning and ending this function (might be data left)
                if curr_recv_timeout_retry < self.recv_timeout_retry:
                    self.log.trace(
                        "Final socket read: {}/{} before quitting.".format(
                            curr_recv_timeout_retry, self.recv_timeout_retry
                        )
                    )
                    curr_recv_timeout_retry += 1
                    continue
                ssig.raise_stop()
                return

            # [2/5] Connection was forcibly closed
            # [Errno 10054] An existing connection was forcibly closed by the remote host
            # [WinError 10054] An existing connection was forcibly closed by the remote host
            # So we will just warn (instead of error our), ignore (don't return here) and
            # let the bottom of this function decide whether we should just quit or
            # reconnect/bind (if the user specified this.
            except socket.error as error:
                # Also ensure we reset the 'byte' value to none so that the below check can kick in.
                byte = None
                self.log.warning("Socket Receive: {}".format(error))

            # [3/5] TODO: Still need to figure out what this error is and when it is thrown
            except AttributeError as error:
                self.log.error("TODO: What happens here?Attribute Receive Error: {}".format(error))
                ssig.raise_stop()
                return

            # We're receiving data again, so let's reset the retry/terminate counter
            # The counter is incremented in 'except socket.timeout' above.
            curr_recv_timeout_retry = 0

            # If we're receiving data from a UDP client
            # we can firstly/finally set its addr/port in order
            # to send data back to it (see send() function)
            if self.options["udp"]:
                self.remote_addr, self.remote_port = addr
                self.log.debug("Client connected: {}:{}".format(self.remote_addr, self.remote_port))

            # [4/5] Upstream (server or client) is gone. Do we reconnect or quit?
            if not byte:
                self.log.trace("Socket: Empty data received or otherwise caught.")

                if self.role == "server":
                    # Yay, we want to continue and allow new clients
                    if self.__reaccept_from_client():
                        self.log.trace("Server can continue, because of --keep-open")
                        continue
                if self.role == "client":
                    # Yay, we want to continue and our client will re-connect upstream again
                    if self.__reconnect_to_server():
                        self.log.trace("Client can continue, because of --reconn")
                        continue

                self.log.warning("Exit. Upstream connection gone. No --keep-open/--reconn set.")
                ssig.raise_stop()
                return

            # [5/5] We have data to process
            data = self.enc.decode(byte)
            self.log.debug(
                "Received {} bytes from {}:{}".format(len(data), self.remote_addr, self.remote_port)
            )
            self.log.trace("Received: {}".format(repr(data)))

            yield data


# -------------------------------------------------------------------------------------------------
# CLASS: NetcatServer
# -------------------------------------------------------------------------------------------------
class NetcatServer(AbstractSocket):
    """Netcat Server implementation."""

    def __init__(self, encoder, host, port, recv_timeout, recv_timeout_retry, options={}):
        """Construct a listening server."""
        super(NetcatServer, self).__init__(
            encoder, "server", recv_timeout, recv_timeout_retry, options
        )

        # Listen on 0.0.0.0 by default
        addr = "0.0.0.0"
        if host is not None:
            addr = self.gethostbyname(host, port, socket.AF_INET)

        # Setup server
        self.create_socket()
        self.bind(addr, port)
        if self.options["udp"]:
            self.conn = self.sock
            self.log.info(
                "Listening on {} (family {}/UDP, port {})".format(addr, socket.AF_INET, port)
            )
        else:
            self.listen()
            self.log.info(
                "Listening on {} (family {}/TCP, port {})".format(addr, socket.AF_INET, port)
            )
            self.accept()


# -------------------------------------------------------------------------------------------------
# CLASS: NetcatClient
# -------------------------------------------------------------------------------------------------
class NetcatClient(AbstractSocket):
    """Netcat Client implementation."""

    def __init__(self, encoder, host, port, recv_timeout, recv_timeout_retry, options={}):
        """Construct a connecting clientt."""
        super(NetcatClient, self).__init__(
            encoder, "client", recv_timeout, recv_timeout_retry, options
        )

        # Setup client
        addr = self.gethostbyname(host, port, socket.AF_INET)
        self.create_socket()
        self.conn = self.sock

        self.remote_addr = addr
        self.remote_port = port
        if not self.options["udp"]:
            if self.connect():
                return
            if self._AbstractSocket__reconnect_to_server():
                return
            sys.exit(1)


# #################################################################################################
# #################################################################################################
# ###
# ###   2 / 3   N E T C A T   P L U G I N S
# ###
# #################################################################################################
# #################################################################################################

# -------------------------------------------------------------------------------------------------
# ABSTRACT CLASS: AbstractNetcatPlugin
# -------------------------------------------------------------------------------------------------
class AbstractNetcatPlugin(ABC):
    """
    Abstract class to for netcat plugins.

    This is a skeleton that defines how the plugins for Netcat should look like.

    The "input_generator" should constantly yield data received from some sort of input
    which could be user input, output from a shell command data from a socket.

    The "input_callback" will apply some sort of action on the data received from a input_generator
    which could be output to stdout, send it to the shell or to a socket.

    "The "input_interrupter" can implement a mechanism to trigger the input_generator to stop
    and return to its parent thread/function. The input_generator must also be implemented
    in a way that it is able to act on the event which the "input_interrupter" emitted.
    """

    __metaclass__ = ABCMeta

    @abstractmethod
    def __init__(self, options={}):
        """
        Set specific options for this plugin.

        Args:
            options (dict):    A dict which allows you to add custom options to your module
        """
        raise NotImplementedError("Should have implemented this")

    @abstractmethod
    def producer(self, ssig):
        """
        Implement a generator function which constantly yields data.

        The data could be from various sources such as: received from a socket,
        received from user input, received from shell command output or anything else.

        Args:
            ssig (StopSignal): A StopSignal instance providing has_stop() and raise_stop() functions
        """
        raise NotImplementedError("Should have implemented this")

    @abstractmethod
    def consumer(self, data):
        """The consumer takes the consumers' output as input and processes it in some form."""
        raise NotImplementedError("Should have implemented this")

    @abstractmethod
    def interrupt(self):
        """Defines an interrupt function which will stop the producer.

        Various producer might call blocking functions and they won't be able to stop themself
        as they hang on that blocking function. This method is triggered from outside and is
        supposed to stop/shutdown the producer.
        If no such interrupt is required, imeplemt it empty."""
        raise NotImplementedError("Should have implemented this")


# -------------------------------------------------------------------------------------------------
# CLASS: NetcatPluginOutput (Module for: user-input -> send -> receive -> output)
# -------------------------------------------------------------------------------------------------
class NetcatPluginOutput(AbstractNetcatPlugin):
    """Implement basic input/output plugin.

    This plugin provides a generator which continuously reads from stdin
    (non-blocking on POSIX and blocking on windows) as well as a
    callback that writes to stdout.
    """

    # Replace '\n' linefeeds (if they exist) with CRLF ('\r\n')?
    crlf = False

    # Non-blocking read from stdin achieved via timeout.
    # Specify timeout in seconds.
    input_timeout = None

    # ------------------------------------------------------------------------------
    # Constructor / Destructor
    # ------------------------------------------------------------------------------
    def __init__(self, options={}):
        """Set specific options for this plugin."""
        super(AbstractNetcatPlugin, self).__init__()
        assert "encoder" in options
        assert "input_timeout" in options
        assert "crlf" in options

        self.log = logging.getLogger(__name__)
        self.enc = options["encoder"]
        if "input_timeout" in options:
            self.input_timeout = options["input_timeout"]
        if "crlf" in options:
            self.crlf = options["crlf"]

    # ------------------------------------------------------------------------------
    # Private Functions
    # ------------------------------------------------------------------------------
    def __use_linefeeds(self, data):
        """Ensure the user input has the desired linefeeds --crlf or not."""
        # No replacement requested or already CRLF
        if not self.crlf or data.endswith("\r\n"):
            return data
        # Replace current newline character with CRLF
        if data.endswith("\n"):
            self.log.debug("Replacing LF with CRLF")
            return data[:-1] + "\r\n"
        # Otherwise just return as it is
        return data

    def __set_input_timeout(self):
        """Throws a catchable BaseException for sys.stdin after timeout (Linux only)."""
        i, o, e = select.select([sys.stdin], [], [], self.input_timeout)
        if not i:
            raise BaseException("timed out")

    # ------------------------------------------------------------------------------
    # Public Functions
    # ------------------------------------------------------------------------------

    def producer(self, ssig):
        """Constantly ask for user input."""
        # https://stackoverflow.com/questions/1450393/#38670261
        # while True: line = sys.stdin.readline() <- reads a whole line (faster)
        # for line in sys.stdin.readlin():        <- reads one byte at a time
        while True:
            if ssig.has_stop():
                self.log.trace("Stop signal acknowledged for reading STDIN-1")
                return
            try:
                # TODO: select() does not work for windows on stdin/stdout
                if os.name != "nt":
                    self.__set_input_timeout()
                line = sys.stdin.readline()
            except BaseException:
                # When using select() with timeout, we don't have any input
                # at this point and simply continue the loop or quit if
                # a terminate request has been made by other threads.
                if ssig.has_stop():
                    self.log.trace("Stop signal acknowledged for reading STDIN-2")
                    return
                continue
            if line:
                self.log.debug("Received {} bytes from STDIN".format(len(line)))
                self.log.trace("Received: {}".format(repr(line)))
                yield self.__use_linefeeds(line)
            # EOF or <Ctrl>+<d>
            else:
                # DO NOT RETURN HERE BLINDLY, THE UPSTREAM CONNECTION MUST GO FIRST!
                if ssig.has_stop():
                    self.log.trace("Stop signal acknowledged for reading STDIN-3")
                    return

    def consumer(self, data):
        """Print received data to stdout."""
        print(data, end="")
        sys.stdout.flush()  # TODO:Is this required? What does this do? Test this!

    def interrupt(self):
        """Empty interrupt."""
        pass


# -------------------------------------------------------------------------------------------------
# CLASS: NetcatPluginCommand (Module for user-input -> send -> execute -> send-back -> output)
# -------------------------------------------------------------------------------------------------
class NetcatPluginCommand(AbstractNetcatPlugin):
    """Implement command execution functionality."""

    executable = None

    # ------------------------------------------------------------------------------
    # Constructor / Destructor
    # ------------------------------------------------------------------------------
    def __init__(self, options={}):
        """Set specific options for this plugin."""
        assert "encoder" in options
        assert "executable" in options

        self.log = logging.getLogger(__name__)
        self.enc = options["encoder"]
        self.executable = options["executable"]
        self.log.debug("Setting '{}' as executable".format(self.executable))

        # Open executable to wait for commands
        env = os.environ.copy()
        try:
            self.p = Popen(
                # TODO: should I also use 'bufsize=1'?
                self.executable,
                stdin=PIPE,
                stdout=PIPE,
                stderr=subprocess.STDOUT,
                shell=False,
                env=env,
            )
        # Python-2 compat (doesn't have FileNotFoundError)
        except OSError:
            self.log.error("Specified executable '{}' not found".format(self.executable))
            sys.exit(1)
        except FileNotFoundError:
            self.log.error("Specified executable '{}' not found".format(self.executable))
            sys.exit(1)

        # Define destructor
        atexit.register(self.__exit__)

    def __exit__(self):
        """Destructor."""
        self.log.trace("Killing executable: {} with pid {}".format(self.executable, self.p.pid))
        self.p.kill()

    def __set_input_timeout(self, timeout=0.1):
        """Throw a TimeOutError Exception for sys.stdin (Linux only)."""
        # select((rlist, wlist, xlist, timeout)
        # rlist: wait until ready for reading
        # wlist: wait until ready for writing
        # xlist: wait for an "exceptional condition"
        i, o, e = select.select([self.p.stdout], [], [], timeout)
        if not i:
            raise BaseException("timed out")

    # ------------------------------------------------------------------------------
    # Public Functions
    # ------------------------------------------------------------------------------
    def interrupt(self):
        """Stop function that can be called externally to close this instance."""
        self.log.trace("[NetcatPluginCommand] subprocess.kill() was raised by input_unterrupter()")
        self.p.kill()

    def producer(self, ssig):
        """Constantly ask for input."""
        while True:
            if ssig.has_stop():
                self.log.trace("Stop signal acknowledged in Command")
                return
            self.log.trace("Reading command output")
            # TODO: non-blocking read does not seem to work or?
            # try:
            # self.__set_input_timeout(timeout=1.5)
            data = self.p.stdout.readline()  # Much better performance than self.p.read(1)
            self.log.trace(data)
            # except BaseException:
            #    if THREAD_TERMINATE:
            #        return
            #    # No input, just check again
            #    #self.p.stdout.flush()
            #    continue
            data = self.enc.decode(data)
            self.log.trace("Command output: {}".format(data))
            if not data:
                self.log.trace("Command output was empty. Exiting loop.")
                break
            yield data

    def consumer(self, data):
        """Send data received to stdin (command input)."""
        data = self.enc.encode(data)
        self.log.trace("Appending to stdin: {}".format(data))
        self.p.stdin.write(data)
        self.p.stdin.flush()


# #################################################################################################
# #################################################################################################
# ###
# ###   3 / 3   N E T C A T   P L U G I N   R U N N E R
# ###
# #################################################################################################
# #################################################################################################

# -------------------------------------------------------------------------------------------------
# CLASS: Runner
# -------------------------------------------------------------------------------------------------
class StopSignal(object):

    __stop = False

    def has_stop(self):
        return self.__stop

    def raise_stop(self):
        self.__stop = True


class Runner(object):
    """Runner class that takes care about putting everything into threads."""

    # Dict of producer/consumer action to run in a thread.
    # Each list item will be run in its own thread
    # {
    #   "name": {
    #     {
    #       "producer": "function",  # A func which yields data
    #       "consumer": "function",  # A callback func to process the data
    #       "signal": "StopSignal",  # An instance providing has_stop() and raise_stop()
    #       "interrupt": "function", # A interrupt func to be called outside the thread to stop it
    #   }
    # }
    __actions = {}

    # Dict of timed function definition to run in a thread.
    # Each list item will be run in its own thread.
    # {
    #   "name": {
    #     {
    #       "action": "function",   # A func which does any action
    #       "intvl": "function",    # The action func will be called at this intervall
    #       "args": "function",     # args for the action func
    #       "kwargs": "function",   # kwargs for the action func
    #       "signal": "StopSignal", # An instance providing has_stop() and raise_stop()
    #   }
    # }
    __timers = {}

    # A dict which holds the threads created from actions.
    # The name is based on the __actions name
    # {"name": "<thread>"}
    __threads = {}

    # ------------------------------------------------------------------------------
    # Constructor / Destructor
    # ------------------------------------------------------------------------------
    def __init__(self):
        """Constructor."""
        self.log = logging.getLogger(__name__)

    # ------------------------------------------------------------------------------
    # Public Functions
    # ------------------------------------------------------------------------------
    def add_action(self, action):
        """
        Enables a function to run threaded by the producer/consumer runner.

        :param str      name:        Name for logging output
        :param function producer:    A generator function which yields data
        :param function consumer:    A callback which consumes data from the generator
        :param function interrupter: A func that signals a stop event to the producer
        """
        assert "name" in action
        assert "producer" in action
        assert "consumer" in action
        assert "signal" in action
        assert "interrupt" in action
        self.__actions[action["name"]] = {
            "name": action["name"],
            "producer": action["producer"],
            "consumer": action["consumer"],
            "signal": action["signal"],
            "interrupt": action["interrupt"],
        }

    def add_timer(self, timer):
        self.__timers[timer["name"]] = {
            "action": timer["action"],
            "intvl": timer["intvl"],
            "args": timer["args"] if "args" in timer else None,
            "kwargs": timer["kwargs"] if "kwargs" in timer else {},
            "signal": timer["signal"],
        }

    def run(self):
        """Run threaded NetCat."""

        def run_action(name, producer, consumer, ssig):
            """
            Receive data (network, user-input, shell-output) and process it (send, output).

            :param str        name:        Name for logging output
            :param function   producer:    A generator function which yields data
            :param function   consumer:    A callback which consumes data from the generator
            :param StopSignal ssig:        Providing has_stop() and raise_stop()
            """
            self.log.trace("[{}] Producer Start".format(name))
            for data in producer(ssig):
                self.log.trace("[{}] Producer received: {}".format(name, repr(data)))
                consumer(data)
            self.log.trace("[{}] Producer Stop".format(name))

        def run_timer(name, action, intvl, ssig, *args, **kwargs):
            """Execute periodic tasks by an optional provided time_action."""
            self.log.trace("[{}] Timer Start (exec every {} sec)".format(name, intvl))
            time_last = int(time.time())
            while True:
                if ssig.has_stop():
                    self.log.trace("Stop signal acknowledged for timer {}".format(name))
                    return
                time_now = int(time.time())
                if time_now > time_last + intvl:
                    self.log.debug("[{}] Executing timed function".format(time_now))
                    # TODO: The following if/else could probably be simplified somehow
                    if args is not None:
                        if kwargs:
                            action(*args, **kwargs)
                        else:
                            action(*args)
                    else:
                        if kwargs:
                            action(**kwargs)
                        else:
                            action()
                    time_last = time_now  # Reset previous time
                time.sleep(0.1)

        # Start available action in a thread
        for key in self.__actions:
            # Create Thread object
            thread = threading.Thread(
                target=run_action,
                name=key,
                args=(
                    key,
                    self.__actions[key]["producer"],
                    self.__actions[key]["consumer"],
                    self.__actions[key]["signal"],
                ),
            )
            thread.daemon = False
            thread.start()
            self.__threads[key] = thread
        # Start available timers in a thread
        for key in self.__timers:
            # Create Thread object
            thread = threading.Thread(
                target=run_timer,
                name=key,
                args=(
                    key,
                    self.__timers[key]["action"],
                    self.__timers[key]["intvl"],
                    self.__timers[key]["signal"],
                    self.__timers[key]["args"],
                ),
                kwargs=self.__timers[key]["kwargs"],
            )
            thread.daemon = False
            thread.start()

        def stop(force):
            """Stop threads."""
            for key in self.__threads:
                if not self.__threads[key].is_alive() or force:
                    self.log.trace("Raise stop signal for {}".format(self.__threads[key].getName()))
                    self.__actions[key]["signal"].raise_stop()
                    self.log.trace("Call interrupt for {}".format(self.__threads[key].getName()))
                    self.__actions[key]["interrupt"]()
                    self.log.trace("Joining {}".format(self.__threads[key].getName()))
                    self.__threads[key].join(timeout=0.1)
            # If all threads have died, exit
            if not all([self.__threads[key].is_alive() for key in self.__threads]) or force:
                if force:
                    sys.exit(1)
                else:
                    sys.exit(0)

        try:
            while True:
                stop(False)
                # Need a timeout to not skyrocket the CPU
                time.sleep(0.1)
        except KeyboardInterrupt:
            print()
            stop(True)


# -------------------------------------------------------------------------------------------------
# COMMAND LINE ARGUMENTS
# -------------------------------------------------------------------------------------------------
def get_version():
    """Return version information."""
    return """%(prog)s: Version %(version)s (%(url)s) by %(author)s""" % (
        {"prog": APPNAME, "version": VERSION, "url": APPREPO, "author": "cytopia"}
    )


def _args_check_port(value):
    """Check arguments for invalid port number."""
    min_port = 1
    max_port = 65535
    intvalue = int(value)

    if intvalue < min_port or intvalue > max_port:
        raise argparse.ArgumentTypeError("%s is an invalid port number" % value)
    return intvalue


def _args_check_forwards(value):
    """Check forward argument (-L/-R) for correct pattern."""
    match = re.search(r"(.+):(.+)", value)
    if match is None or len(match.groups()) != 2:
        raise argparse.ArgumentTypeError("%s is not a valid 'addr:port' format." % value)
    _args_check_port(match.group(2))
    return value


def _args_check_rebind(value):
    """Check rebind argument for correct value."""
    intvalue = int(value)
    if intvalue < 0:
        raise argparse.ArgumentTypeError("must be equal or greater than 0." % value)
    return intvalue


def _args_check_reconn(value):
    """Check reconn argument for correct value."""
    intvalue = int(value)
    if intvalue < 0:
        raise argparse.ArgumentTypeError("must be equal or greater than 0." % value)
    return intvalue


def _args_check_robin_ports(value):
    """Check round-robin argument for comma separated string or range."""
    mcomma = re.search(r"^[0-9]+(,([0-9]+))*$", value)
    mrange = re.search(r"^[0-9]+\-[0-9]+$", value)

    if mcomma is None:
        if mrange is None:
            raise argparse.ArgumentTypeError("%s is not a valid port specifier" % value)

    if mcomma:
        ports = mcomma.group(0).split(",")
        for port in ports:
            _args_check_port(port)
        return ports

    if mrange:
        ranges = mrange.group(0).split("-")
        if int(ranges[0]) >= (int(ranges[1]) + 1):
            raise argparse.ArgumentTypeError(
                "Left side of range must be smaller or equal than right side."
            )
        ports = []
        for port in range(int(ranges[0]), int(ranges[1]) + 1):
            _args_check_port(port)
            ports.append(port)
        return ports


def _args_check_mutually_exclusive(parser, args):
    """Check mutually exclusive arguments."""

    # This is connect mode
    connect_mode = not args.listen and not args.zero and not args.local and not args.remote

    # [MODE] --listen
    if args.listen and (args.zero or args.local or args.remote):
        parser.print_usage()
        print(
            "%s: error: -l/--listen mutually excl. with -z/-zero, -L or -R" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [MODE] --zero
    if args.zero and (args.listen or args.local or args.remote):
        parser.print_usage()
        print(
            "%s: error: -z/--zero mutually excl. with -l/--listen, -L or -R" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [MODE --local
    if args.local and (args.listen or args.zero or args.remote):
        parser.print_usage()
        print(
            "%s: error: -L/--local mutually excl. with -l/--listen, -z/--zero or -R" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [MODE --remote
    if args.remote and (args.listen or args.zero or args.local):
        parser.print_usage()
        print(
            "%s: error: -R/--remote mutually excl. with -l/--listen, -z/--zero or -L" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [MODULE] --exec
    if args.cmd and (args.local or args.remote or args.zero):
        parser.print_usage()
        print(
            "%s: error: -e/--exec mutually excl. with -L, -R or -z/--zero" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [OPTIONS] --udp
    if args.udp and args.zero:
        parser.print_usage()
        print(
            "%s: error: -u/--udp mutually excl. with -z/--zero" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)

    # [ADVANCED] --http
    if args.http and (args.https or args.udp or args.zero):
        parser.print_usage()
        print(
            "%s: error: --http mutually excl. with --https, -u/--udp or -z/--zero" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [ADVANCED] --https
    if args.https and (args.http or args.udp or args.zero):
        parser.print_usage()
        print(
            "%s: error: --https mutually excl. with --http, -z/--udp or -z/--zero" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [ADVANCED] --keep-open
    if args.keep_open and (args.udp):
        parser.print_usage()
        print(
            "%s: error: --keep-open mutually excl. with -u/--udp" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)
    if args.keep_open and not args.listen:
        parser.print_usage()
        print(
            "%s: error: --keep-open only works with -l/--listen" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)

    # [ADVANCED] --rebind
    if args.rebind is not False and not args.listen:
        parser.print_usage()
        print(
            "%s: error: --rebind only works with -l/--listen" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)

    # [ADVANCED] --reconn
    if args.reconn is not False and (args.udp):
        parser.print_usage()
        print(
            "%s: error: --reconn mutually excl. with -u/--udp" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)
    if args.reconn is not False and not (connect_mode or args.zero):
        parser.print_usage()
        print(
            "%s: error: --reconn only works with connect mode or -z/--zero" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # [ADVANCED] --wait
    if args.wait is not False and (args.udp):
        parser.print_usage()
        print(
            "%s: error: -w/--wait mutually excl. with -u/--udp" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)
    if args.wait is not False and not (connect_mode):
        parser.print_usage()
        print(
            "%s: error: -w/--wait only works with connect mode" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)

    # [ADVANCED] --udp-ping-init
    if args.udp_ping_init is not False and (args.listen or args.local):
        parser.print_usage()
        print(
            "%s: error: --udp-ping-init mutually excl. with -l/--listen or -L/--local" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)
    if args.udp_ping_init is not False and not (args.udp):
        parser.print_usage()
        print(
            "%s: error: --udp-ping-init only works with -u/--udp" % (APPNAME), file=sys.stderr,
        )
        sys.exit(1)


def get_args():
    """Retrieve command line arguments."""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawTextHelpFormatter,
        add_help=False,
        usage="""%(prog)s [-Cnuv] [-e cmd] hostname port
       %(prog)s [-Cnuv] [-e cmd] -l [hostname] port
       %(prog)s [-Cnuv] -z hostname port
       %(prog)s [-Cnuv] -L addr:port hostname port
       %(prog)s [-Cnuv] -R addr:port hostname port
       %(prog)s -V, --version
       %(prog)s -h, --help
       """
        % ({"prog": APPNAME}),
        description="""
Enhanced and comptaible Netcat implementation written in Python (2 and 3) with
connect, zero-i/o, listen and forward modes and techniques to detect and evade
firewalls and intrusion detection/prevention systems.

If no mode arguments are specified, pwncat will run in connect mode and act as
a client to connect to a remote endpoint. If the connection to the remote
endoint is lost, pwncat will quit. See advanced options for how to automatically
reconnect.""",
    )

    positional = parser.add_argument_group("positional arguments")
    mode = parser.add_argument_group("mode arguments")
    optional = parser.add_argument_group("optional arguments")
    advanced = parser.add_argument_group("advanced arguments")
    misc = parser.add_argument_group("misc arguments")

    positional.add_argument(
        "hostname", nargs="?", type=str, help="Address to listen, forward or connect to"
    )
    positional.add_argument(
        "port", type=_args_check_port, help="Port to listen, forward or connect to"
    )

    mode.add_argument(
        "-l",
        "--listen",
        action="store_true",
        default=False,
        help="""[Listen mode]:
Start a server and listen for incoming connections.
If using TCP and a connected client disconnects or the
connection is interrupted otherwise, the server will
quit. See -k/--keep-open to change this behaviour.

""",
    )
    mode.add_argument(
        "-z",
        "--zero",
        action="store_true",
        default=False,
        help="""[Zero-I/0 mode]:
Connect to a remote endpoint and report status only.
Used for port scanning.

""",
    )
    mode.add_argument(
        "-L",
        "--local",
        metavar="addr:port",
        default=False,
        type=_args_check_forwards,
        help="""[Local forward mode]:
This mode will start a server and a client internally.
The internal server will listen locally on specified
hostname/port (positional arguments). Same as with -l.
The server will then forward traffic to the internal
client which connects to another server specified by
address given via -L/--local addr:port.
(I.e.: proxies a remote service to a local address)

"""
        % ({"prog": APPNAME}),
    )
    mode.add_argument(
        "-R",
        "--remote",
        metavar="addr:port",
        default=False,
        type=_args_check_forwards,
        help="""[Remote forward mode]:
This mode will start two clients internally. One is
connecting to the target and one is connecting to
another pwncat/netcat server you have started some-
where. Once connected, it will then proxy traffic
between you and the target.
This mode should be applied on machines that block
incoming traffic and only allow outbound.
The connection to your listening server is given by
-R/--remote addr:port and the connection to the
target machine via the positional arguments.
"""
        % ({"prog": APPNAME}),
    )
    optional.add_argument(
        "-e",
        "--exec",
        metavar="cmd",
        dest="cmd",
        default=False,
        type=str,
        help="Execute shell command. Only for connect or listen mode.",
    )
    optional.add_argument(
        "-C",
        "--crlf",
        action="store_true",
        default=False,
        help="Replace LF with CRLF from stdin (default: don't)",
    )
    optional.add_argument(
        "-n", "--nodns", action="store_true", default=False, help="Do not resolve DNS.",
    )
    optional.add_argument(
        "-u",
        "--udp",
        action="store_true",
        default=False,
        help="Use UDP for the connection instead of TCP.",
    )
    optional.add_argument(
        "-v",
        "--verbose",
        action="count",
        default=0,
        help="""Be verbose and print info to stderr. Use -v, -vv, -vvv
or -vvvv for more verbosity. The server performance will
decrease drastically if you use more than three times.""",
    )

    advanced.add_argument(
        "--http",
        action="store_true",
        default=False,
        help="""Connect / Listen / Local forward mode (TCP only):
Hide traffic in http packets to fool Firewalls/IDS/IPS.

""",
    )
    advanced.add_argument(
        "--https",
        action="store_true",
        default=False,
        help="""Connect / Listen / Local forward mode (TCP only):
Hide traffic in https packets to fool Firewalls/IDS/IPS.

""",
    )
    advanced.add_argument(
        "-k",
        "--keep-open",
        action="store_true",
        default=False,
        help="""Listen mode (TCP only):
Re-accept new clients in listen mode after a client has
disconnected or the connection is unterrupted otherwise.
(default: server will quit after connection is gone)

""",
    )
    advanced.add_argument(
        "--rebind",
        metavar="x",
        default=False,
        type=_args_check_rebind,
        help="""Listen mode (TCP and UDP):
If the server is unable to bind, it will re-initialize
itself x many times before giving up. Use 0 to re-init
endlessly. (default: fail after first unsuccessful try).

""",
    )
    advanced.add_argument(
        "--rebind-wait",
        metavar="s",
        default=1,
        type=int,
        help="""Listen mode (TCP and UDP):
Wait x seconds between re-initialization. (default: 1)

""",
    )
    advanced.add_argument(
        "--rebind-robin",
        metavar="port",
        default=[],
        type=_args_check_robin_ports,
        help="""Listen mode (TCP and UDP):
If the server is unable to initialize (e.g: cannot bind
and --rebind is specified, it it will shuffle ports in
round-robin mode to bind to. Use comma separated string
such as '80,81,82' or a range of ports '80-100'.
Set --rebind to at least the number of ports to probe +1
This option requires --rebind to be specified.

""",
    )
    advanced.add_argument(
        "--reconn",
        metavar="x",
        default=False,
        type=_args_check_reconn,
        help="""Connect mode / Zero-I/O mode (TCP only):
If the remote server is not reachable or the connection
is interrupted, the client will connect again x many
times before giving up. Use 0 to retry endlessly.
(default: quit if the remote is not available or the
connection was interrupted)
This might be handy for stable TCP reverse shells ;-)

""",
    )
    advanced.add_argument(
        "--reconn-wait",
        metavar="s",
        default=1,
        type=int,
        help="""Connect mode / Zero-I/O mode (TCP only):
Wait x seconds between re-connects. (default: 1)

""",
    )
    advanced.add_argument(
        "--reconn-robin",
        metavar="port",
        default=[],
        type=_args_check_robin_ports,
        help="""Connect mode / Zero-I/O mode (TCP only):
If the remote server is not reachable or the connection
is interrupted and --reconn is specified, the client
will shuffle ports in round-robin mode to connect to.
Use comma separated string such as '80,81,82' or a range
of ports '80-100'.
Set --reconn to at least the number of ports to probe +1
This helps reverse shell to evade intrusiona prevention
systems that will cut your connection and block the
outbound port.
This is also useful in Connect or Zero-I/O mode to
figure out what outbound ports are allowed.

""",
    )
    advanced.add_argument(
        "-w",
        "--wait",
        metavar="s",
        default=False,
        type=int,
        help="""Connect mode (TCP only):
If a connection and stdin are idle for more than s sec,
then the connection is silently closed and the client
will exit. (default: wait forever).
Note: if --reconn is specified, the connection will be
re-opened.

""",
    )
    advanced.add_argument(
        "--udp-ping-init",
        action="store_true",
        default=False,
        help="""Connect mode / Zero-I/O mode (UDP only):
UDP is a stateless protocol unlike TCP, so no hand-
shake communication takes place and the client just
sends data to a server without being "accepted" by
the server first.
This means a server waiting for an UDP client to
connect to, is unable to send any data to the client,
before the client hasn't send data first. The server
simply doesn't know the IP address before an initial
connect.
The --udp-ping-init option instructs the client to send
one single initial ping packet to the server, so that it
is able to talk to the client.
This is the only way to make a UDP reverse shell work.
See --udp-ping-word for what char/string to send as
initial ping packet (default: '\\0')

""",
    )
    advanced.add_argument(
        "--udp-ping-intvl",
        metavar="s",
        default=False,
        type=int,
        help="""Connect mode / Zero-I/O mode (UDP only):
Instruct the UDP client to send ping intervalls every
s seconds. This allows you to restart your UDP server
and just wait for the client to report back in.
This might be handy for stable UDP reverse shells ;-)
See --udp-ping-word for what char/string to send as
initial ping packet (default: '\\0')

""",
    )
    advanced.add_argument(
        "--udp-ping-word",
        metavar="str",
        default="\0",
        type=str,
        help="""Connect mode / Zero-I/O mode (UDP only):
Change the default character '\\0' to use for upd ping.
Single character or strings are supported.

""",
    )
    advanced.add_argument(
        "--udp-ping-robin",
        metavar="port",
        default=[],
        type=_args_check_robin_ports,
        help="""Zero-I/O mode (UDP only):
Instruct the UDP client to shuffle the specified ports
in round-robin mode for a remote server to ping.
This might be handy to scan outbound allowed ports.
Use --udp-ping-intvl 0 to be faster.

""",
    )
    advanced.add_argument(
        "--safe-word",
        metavar="str",
        default=False,
        type=str,
        help="""All modes:
If %(prog)s is started with this argument, it will shut
down as soon as it receives the specified string. The
--keep-open (server) or --reconn (client) options will
be ignored and it won't listen again or reconnect to you.
Use a very unique string to not have it shut down
accidentally by other input.
"""
        % ({"prog": APPNAME}),
    )
    misc.add_argument("-h", "--help", action="help", help="Show this help message and exit")
    misc.add_argument(
        "-V",
        "--version",
        action="version",
        version=get_version(),
        help="Show version information and exit",
    )

    # Retrieve arguments
    args = parser.parse_args()

    # Check mutually exclive arguments
    _args_check_mutually_exclusive(parser, args)

    # Connect mode and Zero-I/O mode require hostname and port to be set
    connect_mode = not (args.listen or args.zero or args.local or args.remote)
    if (connect_mode or args.zero) and not (args.hostname):
        parser.print_usage()
        print(
            "%s: error: the following arguments are required: hostname" % (APPNAME),
            file=sys.stderr,
        )
        sys.exit(1)

    # Deny unimplemented modes
    if args.http or args.https or args.rebind or args.wait or args.udp_ping_init or args.safe_word:
        print("Unimplemented options", file=sys.stderr)
        sys.exit(1)

    return args


# -------------------------------------------------------------------------------------------------
# CUSTOM FUNCTIONS
# -------------------------------------------------------------------------------------------------
def logtrace(self, message, *args, **kws):
    """Set custom log level for TRACE."""
    if self.isEnabledFor(LOGLEVEL_TRACE_NUM):
        self._log(LOGLEVEL_TRACE_NUM, message, args, **kws)


# -------------------------------------------------------------------------------------------------
# MAIN ENTRYPOINT
# -------------------------------------------------------------------------------------------------
def main():
    """Run the program."""
    args = get_args()
    host = args.hostname
    port = args.port

    # Set netcat options
    net_opts = {
        "bufsize": 8192,
        "backlog": 0,
        "nodns": args.nodns,
        "udp": args.udp,
        "http": args.http,
        "https": args.https,
        "keep_open": args.keep_open,
        "rebind": True if (type(args.rebind) is int and args.rebind == 0) else args.rebind,
        "reconn": True if (type(args.reconn) is int and args.reconn == 0) else args.reconn,
        "rebind_wait": args.rebind_wait,
        "reconn_wait": args.reconn_wait,
        "rebind_robin": args.rebind_robin,
        "reconn_robin": args.reconn_robin,
        "wait": args.wait,
        "udp_ping_init": args.udp_ping_init,
        "udp_ping_intvl": args.udp_ping_intvl,
        "udp_ping_word": args.udp_ping_word,
        "udp_ping_robing": args.udp_ping_robin,
        "safe_word": args.safe_word,
    }

    # Initialize logger
    if args.verbose == 0:
        loglevel = logging.ERROR
    elif args.verbose == 1:
        loglevel = logging.WARNING
    elif args.verbose == 2:
        loglevel = logging.INFO
    elif args.verbose == 3:
        loglevel = logging.DEBUG
    else:
        loglevel = LOGLEVEL_TRACE_NUM

    logging.addLevelName(LOGLEVEL_TRACE_NUM, "TRACE")
    logging.Logger.trace = logtrace
    logformat = "%(levelname)s %(message)s"
    if args.verbose > 2:
        logformat = "%(levelname)s [%(threadName)s]: %(message)s"
    if args.verbose > 3:
        logformat = "%(levelname)s [%(threadName)s] %(lineno)d:%(funcName)s(): %(message)s"
    logging.basicConfig(format=logformat, level=loglevel)

    # Initialize encoder
    encoder = StringEncoder()

    # Use command modulde
    if args.cmd:
        module_opts = {
            "encoder": encoder,
            "executable": args.cmd,
        }
        mod = NetcatPluginCommand(module_opts)
    # Use output module
    else:
        module_opts = {
            "encoder": encoder,
            "crlf": args.crlf,
            "input_timeout": TIMEOUT_READ_STDIN,
        }
        mod = NetcatPluginOutput(module_opts)

    # Run local port-forward
    # -> listen locally and forward traffic to remote (connect)
    if args.local:
        ssig = StopSignal()
        # TODO: Make the listen address optional!
        srv_opts = net_opts.copy()
        srv_opts["reconn"] = True
        srv_opts["keep_open"] = True
        srv_opts["reconn_wait"] = 0
        lhost = args.local.split(":")[0]
        lport = int(args.local.split(":")[1])
        # Create listen and client instances
        net_srv = NetcatServer(
            encoder, lhost, lport, TIMEOUT_RECV_SOCKET, TIMEOUT_RECV_SOCKET_RETRY, srv_opts
        )
        net_cli = NetcatClient(
            encoder, host, port, TIMEOUT_RECV_SOCKET, TIMEOUT_RECV_SOCKET_RETRY, srv_opts
        )
        # Create Runner
        run = Runner()
        run.add_action(
            {
                "name": "TRANSMIT",
                "producer": net_srv.receive,  # USER sends data to PC-SERVER
                "consumer": net_cli.send,  # Data parsed on to PC-CLIENT to send to TARGET
                "signal": ssig,
                "interrupt": object,
            }
        )
        run.add_action(
            {
                "name": "RECEIVE",
                "producer": net_cli.receive,  # Data comes back from TARGET to PC-CLIENT
                "consumer": net_srv.send,  # Data parsed on to PC-SERVER to back send to USER
                "signal": ssig,
                "interrupt": object,
            }
        )
        run.run()
    # Run remote port-forward
    # -> connect to client, connect to target and proxy traffic in between.
    if args.remote:
        ssig = StopSignal()
        # TODO: Make the listen address optional!
        srv_opts = net_opts.copy()
        srv_opts["reconn"] = True
        srv_opts["reconn_wait"] = 0.2
        lhost = args.remote.split(":")[0]
        lport = int(args.remote.split(":")[1])
        # Create local and remote client
        net_cli_l = NetcatClient(
            encoder, lhost, lport, TIMEOUT_RECV_SOCKET, TIMEOUT_RECV_SOCKET_RETRY, srv_opts
        )
        net_cli_r = NetcatClient(
            encoder, host, port, TIMEOUT_RECV_SOCKET, TIMEOUT_RECV_SOCKET_RETRY, srv_opts
        )
        # Create Runner
        run = Runner()
        run.add_action(
            {
                "name": "TRANSMIT",
                "producer": net_cli_l.receive,  # USER sends data to PC-SERVER
                "consumer": net_cli_r.send,  # Data parsed on to PC-CLIENT to send to TARGET
                "signal": ssig,
                "interrupt": object,
            }
        )
        run.add_action(
            {
                "name": "RECEIVE",
                "producer": net_cli_r.receive,  # Data comes back from TARGET to PC-CLIENT
                "consumer": net_cli_l.send,  # Data parsed on to PC-SERVER to back send to USER
                "signal": ssig,
                "interrupt": object,
            }
        )
        run.run()
    # Run server
    if args.listen:
        ssig = StopSignal()
        net = NetcatServer(
            encoder, host, port, TIMEOUT_RECV_SOCKET, TIMEOUT_RECV_SOCKET_RETRY, net_opts
        )
        run = Runner()
        run.add_action(
            {
                "name": "RECV",
                "producer": net.receive,
                "consumer": mod.consumer,
                "signal": ssig,
                "interrupt": mod.interrupt,  # Also force the producer to stop on net error
            }
        )
        run.add_action(
            {
                "name": "STDIN",
                "producer": mod.producer,
                "consumer": net.send,
                "signal": ssig,
                "interrupt": mod.interrupt,  # Externally stop the produer itself
            }
        )
        run.run()

    # Run client
    else:
        ssig = StopSignal()
        net = NetcatClient(
            encoder, host, port, TIMEOUT_RECV_SOCKET, TIMEOUT_RECV_SOCKET_RETRY, net_opts
        )
        run = Runner()
        run.add_action(
            {
                "name": "RECV",
                "producer": net.receive,
                "consumer": mod.consumer,
                "signal": ssig,
                "interrupt": mod.interrupt,  # Also force the producer to stop on net error
            }
        )
        run.add_action(
            {
                "name": "STDIN",
                "producer": mod.producer,
                "consumer": net.send,
                "signal": ssig,
                "interrupt": mod.interrupt,  # Externally stop the produer itself
            }
        )
        if type(args.udp_ping_intvl) is int and args.udp_ping_intvl > 0:
            run.add_timer(
                {
                    "name": "PING",
                    "action": net.send,
                    "intvl": args.udp_ping_intvl,
                    "args": ("\x00"),
                    "signal": ssig,
                }
            )
        run.run()


if __name__ == "__main__":
    # Catch Ctrl+c and exit without error message
    try:
        main()
    except KeyboardInterrupt:
        print()
        sys.exit(1)
