#!/usr/bin/env python3
# author: deadc0de6
#
# A simple python echo server with TCP and UDP support
# able to handle different responses:
#   - echo
#   - payload from arg
#   - payload from file
#
# Copyright (C) 2015 deadc0de6
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#

import sys
import socket
import select
import os
import _thread
from docopt import docopt # sudo pip install docopt

VERSION = 0.2
VERBOSE = False
BUFSZ = 4096
BACKLOG = 10
TIMEOUT = 3 # seconds
USAGE = '''
{0}

Usage:
    {0} tcp <port> [-Vi] [--listen=(local|global)] [--payload=<pld>|--fpayload=<path>]
    {0} udp <port> [-V] [--listen=(local|global)] [--payload=<pld>|--fpayload=<path>]
    {0} --version
    {0} --help

Options:
    -p --payload=<pld>     Payload to send (echo if not set).
    -f --fpayload=<path>   Payload from file to send (echo if not set).
    -l --listen=value      Listen on 127.0.0.1 or 0.0.0.0 [default: local].
    -i --first             Send payload once remote connects [default: False].
    -V --verbose           Be verbose.
    -h --help              Show help.
    -v --version           Show version.
'''

def verbose(header, message):
    sys.stderr.write('[%s] %s\n' % (header, message))

# ############################################################################
# TCP
# ############################################################################
def tcp_answer(connection, client, payload=None):
    if VERBOSE:
        verbose('%s:%u' % (client[0], client[1]), 'connection opened')
    try:
        if FIRST:
            connection.sendall(payload.encode())
        while True:
            R, W, E = select.select([connection], [], [connection], TIMEOUT)
            if R == [] and W == [] and E == []:
                # timed out
                break
            for s in R:
                data = s.recv(BUFSZ)
                if data:
                    if VERBOSE:
                        verbose('%s:%u' % (client[0], client[1]), 'data received')
                    if payload:
                        s.sendall(payload.encode())
                    else:
                        s.sendall(data)
                    if VERBOSE:
                        verbose('%s:%u' % (client[0], client[1]), 'data sent')
            for s in E:
                sys.stderr.write('errors')
                break
    finally:
        connection.close()
    if VERBOSE:
        verbose('%s:%u' % (client[0], client[1]), 'connection closed')

def tcp_echo(port, address, payload=None):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    addr = (address, port)
    s.bind(addr)
    s.listen(BACKLOG)
    print('listening on TCP/{}:{}'.format(ADDR, str(PORT)))
    try:
        while True:
            conn, client = s.accept()
            _thread.start_new_thread(tcp_answer, (conn, client, payload,))
    except KeyboardInterrupt:
        sys.stderr.write('aborted\n')
        pass
    finally:
        s.close()

# ############################################################################
# UDP
# ############################################################################
def udp_echo(port, address, payload=None):
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    addr = (address, port)
    s.bind(addr)
    print('listening on UDP/{}:{}'.format(ADDR, str(PORT)))
    try:
        while True:
            data, remote = s.recvfrom(BUFSZ)
            if VERBOSE:
                verbose('%s:%u' % (remote[0], remote[1]), 'data received')
            if data:
                if payload:
                    s.sendto(payload.encode(), remote)
                else:
                    s.sendto(data, remote)
                if VERBOSE:
                    verbose('%s:%u' % (remote[0], remote[1]), 'data sent')
    except KeyboardInterrupt:
        sys.stderr.write('aborted\n')
        pass
    finally:
        s.close()

# ############################################################################
# Entry point
# ############################################################################
args = docopt(USAGE, version=VERSION)
# parse args
VERBOSE = args['--verbose']
PORT = int(args['<port>'])
ADDR = '127.0.0.1'
FIRST = args['--first']
if args['global']:
  ADDR = '0.0.0.0'

PLD = None
if args['--payload']:
    PLD = args['--payload']
elif args['--fpayload']:
    path = args['--fpayload']
    if not os.path.exists(path):
        print('path \"%s\" does not exist !' % (path))
        sys.exit(0)
    f = open(path, 'r')
    PLD = f.read()
    f.close()

if FIRST and not PLD:
    print('you need to provide a payload with --first')
    sys.exit(1)

# listen and answer
print('... abort with CTRL+c ...')
if args['tcp']:
    tcp_echo(PORT, ADDR, payload=PLD)
elif args['udp']:
    udp_echo(PORT, ADDR, payload=PLD)

