#!/usr/bin/python
# copyright Greg Novak (novak@ucolick.org) 2004
# Released under the GPL, available from http://www.gnu.org

import getopt,sys,socket,tempfile,os

def usage():
    print """
Read or write things from sockets.

Usage: socket [-hl] [-d <directory>]
              [-e <command> [-L] | -p <command> | -f <file> ]
              host port
     -h : help
     -l : listen (default: print to stdout)
     -e : execute <command> on data
     -L : Leave files after <command> finishes
     -p : pipe data to <command>
     -d : use directory <dir>
     -f : write data to <file>
        : to get sequential files, include a %d
        : otherwise, files will be overwritten
"""

try:
    opts, args = getopt.getopt(sys.argv[1:], "hlLe:p:f:d:")
except getopt.GetoptError:
    # print help information and exit:
    usage()
    sys.exit(2)
#Global vars
buffer=1024

#Global switches for options
listen=False
execute=False
pipe=False
file=False
overwriting_file=False
LeaveFiles=False
command=""

for o, a in opts:
    if o == "-l":
        listen = True
    if o == "-h":
        usage()
        sys.exit()
    if o == "-d":
        os.chdir(a)
    if o == "-L":
        LeaveFiles=True
    if o == "-e":
        command = a
        execute = True
    if o == "-p":
        command = a
        pipe = True
    if o == "-f":
        command = a
        if command.find("%d") > 0: file = True
        else: overwriting_file=True
        
# a little error checking
if (execute and pipe) or (execute and file) or (pipe and file):
    usage()
    sys.exit()

# finally, grab the port number
host = args[0]
port = int(args[1])

if listen:
    # Server socket -- receive shit
    s=socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind((host, port))    
    s.listen(5)
    filenum = 1;
    
    while True:
        try:
            (cs, addr) = s.accept()
        except KeyboardInterrupt:
            print "Finished!"
            s.close()
            sys.exit()
            
        # Get the file ready
        if execute:
            filename = tempfile.mktemp()
            f=open(filename, 'w')
        elif pipe:
            f=os.popen(command, 'w')
        elif file:
            f=open(command % filenum, 'w')
        elif overwriting_file:
            f=open(command, 'w')
        else:
            f = sys.stdout

        # Read the socket
        chunk = cs.recv(buffer);
        while len(chunk) > 0:
            f.write(chunk)
            chunk = cs.recv(buffer);
        cs.close()

        # cleanup
        filenum += 1
        if execute or pipe or file or overwriting_file:
            f.close()
            
        if execute:
            os.system(command + " " + filename);
            if not LeaveFiles:
                os.unlink(filename)

else:    
    # client socket -- send shit
    s=socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((host, port))
    ln = sys.stdin.readline()
    while (ln != ''):
        sent = 0
        while sent < len(ln):
            chunk = s.send(ln[sent:])
            sent += chunk
        ln = sys.stdin.readline()
    s.close()
