tcp over file python

file_comm.py

# from https://labs.f-secure.com/tools/tcp-over-file-tunnel/
import sys
import socket
import threading
import base64
import time
import binascii
import mutex
import signal

# Usage
if (len(sys.argv) != 6):
    print "Usage: %s <mode> <ip_address> <tcp_port> <read_file> <write_file>\n" % sys.argv[0]
    print "\tmode = 1 (client - connects to <tcp_port> on <ip_address>)"
    print "\tmode = 2 (server - listens on <tcp_port> on <ip_address>)"
    sys.exit(1)


# Assign program arguments
mode = int(sys.argv[1])
address = sys.argv[2]
port = int(sys.argv[3])
read_filename = sys.argv[4]
write_filename = sys.argv[5]


# Setup useful variables 
lock = mutex.mutex()
buffered_data = {}

# Clear files
fp = open(read_filename, "w")
fp.close()
fp = open(write_filename, "w")
fp.close()


# Open read and write files
read_fp = open(read_filename, "r")
write_fp = open(write_filename, "w")


# interrupt signal handler
def sig_handler(signum, frame):
    print "Signal caught, exiting..."
    sys.exit(1)


# register interrupt signal with handle
signal.signal(signal.SIGABRT, sig_handler)
print "[*] Prese Ctrl-Break to quit...\n"


# loop listening for new socket connections
def connection_accepter():
    global lock
    
    serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    serversocket.bind((address, port))
    next_conn_id = 1

    while True:
        serversocket.listen(5)
        (s, clientaddress) = serversocket.accept()
        conn_id = next_conn_id
        next_conn_id += 1
        print "[*] Connection received (ID = %d) from %s:%d" % (conn_id, clientaddress[0], clientaddress[1])  

        while(lock.testandset() == False):
           pass

        write_fp.write("%d #CONNECT#$" % conn_id)
        write_fp.flush()
        buffered_data[conn_id] = []

        lock.unlock()
        
        t = threading.Thread(target = socket_reader_thread, name="t%d" % conn_id, args=[conn_id, s])
        t.start()
        t = threading.Thread(target = socket_writer_thread, name="t%d" % conn_id, args=[conn_id, s])
        t.start()

        


# read from file and process connections IDs + data
def file_reader():
    packet_buffer = ""
    while True:
        packet = read_fp.read(1024)
        if (packet != ''):
            packet_buffer += packet
            while True:
                (part_before, part_sep, part_after) = packet_buffer.partition("$")
                if (part_sep == ''):
                    break
                process_packet(part_before)
                packet_buffer = part_after
        else:
            time.sleep(0.001)    


# Read from socket/Write to file
def socket_reader_thread(conn_id, s):
    global lock
    
    while True:
        if (buffered_data.has_key(conn_id)):
            try:
                data = s.recv(768)
            except socket.error, (errno, strerror):
                print "[*] (ID = %d): %d %s" % (conn_id, errno, strerror)

                while(lock.testandset() == False):
                    pass

                write_fp.write("%d #DISCONNECT#$" % (conn_id))
                write_fp.flush()
                del buffered_data[conn_id]
                s.close()

                lock.unlock()
                break
                
            while(lock.testandset() == False):
                pass
            
            if (data != ''):
                encoded_data = base64.b64encode(data)
                write_fp.write("%d %s$" % (conn_id, encoded_data))
                write_fp.flush()

                #print "data read from socket (%d)" % len(encoded_data)
                #print "data read from socket (%s)" % encoded_data
            lock.unlock()
        else:
            s.close()
            break


# Read from connection buffer/Write to socket
def socket_writer_thread(conn_id, s):
    global buffered_data, lock

    while True:
        if (buffered_data.has_key(conn_id)):
            if (len(buffered_data[conn_id]) > 0):
                while(lock.testandset() == False):
                    pass

                try:
                    data = buffered_data[conn_id].pop(0)
                except KeyError, (errno):
                    lock.unlock()
                    break
                
                lock.unlock()
                
                s.send(data)

                #print "data read from socket (%d)" % len(data)
                #print "data read from socket (%s)" % data
            else:
                time.sleep(0.001)
        else:
            break


# Process file packet
def process_packet(packet):
    global lock
    
    (conn_id, data) = packet.split(" ")
    conn_id = int(conn_id)

    while(lock.testandset() == False):
        pass

    if (data == "#CONNECT#"):
        print "[*] Connection request received (ID = %d). Connecting to %s on port %d" % (conn_id, address, port)    
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((address, port))
            
        buffered_data[conn_id] = []
        t = threading.Thread(target = socket_reader_thread, name="r%d" % conn_id, args=[conn_id, s])
        t.start()
        t = threading.Thread(target = socket_writer_thread, name="w%d" % conn_id, args=[conn_id, s])
        t.start()
    elif (data == "#DISCONNECT#"):
        print "[*] Disconnect request received (ID = %d). Connection terminated " % (conn_id)
        del buffered_data[conn_id]
    else:
        decoded_data = base64.b64decode(data)
        try:
            buffered_data[conn_id].append(decoded_data)
        except KeyError, (errno):
            pass
        
        #print "data written to socket (%d)" % len(data)
        #print "data written to socket (%s)" % data
    
    lock.unlock()



# Mode specific configuration
if (mode == 1):        
    pass
elif (mode == 2):
    t1 = threading.Thread(target = connection_accepter, name="connection accepter", args=[])
    t1.start()
else:
    print "Error: invalid mode\n"
    sys.exit(1)


# Fire off file reader thread
t1 = threading.Thread(target = file_reader, name="file reader", args=[])
t1.start()

while True:
    time.sleep(0.1)