mtftp.py - TFTP server and client

'''
mtftp.py

version 1.0

Description
   tftp client and server as described in RFC 1350, more or less

History
Version 0.0 September 2015;

   initial implementation.

Version 1.0 October 6th 2015

   client and server implemented



Author: Guillermo Gomez (a.k.a Momo)

'''

import argparse
import os
import dpkt
import socket, random
import sys, time

# Opcodes
OP_RRQ     = 1    # read request
OP_WRQ     = 2    # write request
OP_DATA    = 3    # data packet
OP_ACK     = 4    # acknowledgment
OP_ERR     = 5    # error code

# Error codes
EUNDEF     = 0    # not defined
ENOTFOUND  = 1    # file not found
EACCESS    = 2    # access violation
ENOSPACE   = 3    # disk full or allocation exceeded
EBADOP     = 4    # illegal TFTP operation
EBADID     = 5    # unknown transfer ID
EEXISTS    = 6    # file already exists
ENOUSER    = 7    # no such user

tftp_err_code = ['EUNDEF','ENOTFOUND','EACCESS','ENOSPACE','EBADOP','EBADID','EEXISTS','ENOUSER']

tftp_err_msg = ['not defined',\
                'file not found',\
                'access violation',\
                'disk full or allocation exceeded',\
                'illegal TFTP operation',\
                'unknown transfer ID',\
                'file already exists',\
                'no such user']

def Send_ERROR(sckt, peer_add, error_code):
    ERR_pkt = dpkt.tftp.TFTP(opcode = 5,\
                             errcode = error_code,\
                             errmsg = tftp_err_msg[error_code])
    return sckt.sendto(ERR_pkt.pack(), peer_add)

def Send_ACK(sckt, peer_add, block_num):
    ACK_pkt = dpkt.tftp.TFTP(opcode = 4,\
                             block = block_num)
    return sckt.sendto(ACK_pkt.pack(), peer_add)

def Send_RRQ(sckt, peer_add, file_name):
    RRQ_pkt = dpkt.tftp.TFTP(opcode = 1,\
                             filename = file_name,\
                             mode = 'octet')
    return sckt.sendto(RRQ_pkt.pack(), peer_add)

def Send_WRQ(sckt, peer_add, file_name):
    WRQ_pkt = dpkt.tftp.TFTP(opcode = 2,\
                             filename = file_name,\
                             mode = 'octet')
    return sckt.sendto(WRQ_pkt.pack(), peer_add)

def Send_DATA(sckt, peer_add, block_num, block_data):
    DATA_pkt = dpkt.tftp.TFTP(opcode = 3,\
                             block = block_num,\
                             data = block_data)
    return sckt.sendto(DATA_pkt.pack(), peer_add)

def Read_File(srv_add, file_name, max_ret=3): # Read file from server

    # srv_add is a tuple ("ip.addr", TID )
    # returns successful condition, result message, received byte number and time elapsed
       
    # init variables
    block_num = 1 # first transmited block number
    ret_num = 0
    bytes_received = 0
    ret_cnt = max_ret
    result_msg = ''
    last_pkt = False
    init_time = time.time()

    try:
        local_file = open (file_name , "wb" )
    except IOError as message:
        result_msg = str(message)
        print '\n' + result_msg
        return False, result_msg, 0, 0

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client_socket.settimeout(3)
    Send_RRQ (client_socket, srv_add, file_name)
    print ('Reading file "{}" from server "{}"'.format(file_name, srv_add[0]))

    while (ret_cnt > 0):
        try:
            data, source = client_socket.recvfrom(1024)
            if block_num == 1 and (srv_add[0] == source[0]): # remember server port
                srv_add = source
            if srv_add == source: # thats my flow
                tftp_pkt = dpkt.tftp.TFTP(data)
                if tftp_pkt.opcode == dpkt.tftp.OP_DATA:
                    if len(tftp_pkt.data) > 512:
                        result_msg = tftp_err_msg[EBADOP]
                        Send_ERROR(client_socket, srv_add, EBADOP)
                        break
                    elif len(tftp_pkt.data) < 512: # last block
                        last_pkt = True
                        result_msg = 'Transfer successful'
                    if tftp_pkt.block[0] == block_num: # new_block
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                        local_file.write(tftp_pkt.data)
                        bytes_received += len(tftp_pkt.data)
                        Send_ACK(client_socket, srv_add, tftp_pkt.block[0])
                        block_num += 1
                        if not last_pkt:
                            ret_cnt = max_ret
                        else:
                            ret_cnt = 1 # for dallying
                    elif tftp_pkt.block[0] == (block_num -1): # retransmission
                        ret_num += 1
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                        ret_cnt -= 1
                        Send_ACK(client_socket, srv_add, tftp_pkt.block[0])
                    else:
                        result_msg = tftp_err_msg[EBADOP]
                        Send_ERROR(client_socket, srv_add, EBADOP)
                        break
                elif tftp_pkt.opcode == dpkt.tftp.OP_ERR:
                    result_msg = tftp_pkt.errmsg
                    break
                else:
                    result_msg = tftp_err_msg[EBADOP]
                    Send_ERROR(client_socket, srv_add, EBADOP)
                    break
            else: # source != srv_add , that is not my flow
                result_msg = tftp_err_msg[EBADID]
                Send_ERROR(client_socket, source, EBADID)
                continue
        except socket.timeout:
            if last_pkt:
                break
            elif block_num == 1:
                Send_RRQ(client_socket, srv_add, file_name)
            else:
                Send_ACK(client_socket, srv_add, (block_num-1))
                ret_num += 1
                print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
            ret_cnt -= 1
        except socket.error as message:
            result_msg = str(message)
            Send_ERROR(client_socket, srv_add, EUNDEF)
            break
        except IOError as message:
            result_msg = str(message)
            Send_ERROR(client_socket, srv_add, ENOSPACE)
            break
    client_socket.close()
    local_file.close()
    total_time = time.time()-init_time
    print '\n' + result_msg
    print ('Received {} bytes in {:>.3f} seconds. Data Rate {:>.0f} B/s'.format(\
        bytes_received, total_time, bytes_received/total_time))
    return last_pkt, result_msg, bytes_received, total_time

def Write_File(srv_add, file_name, max_ret=3): # Write file to server

    # TID (a.k.a port)
    # srv_add is a tuple ("ip.addr", TID )
    # returns successful condition, result message, received byte number and time elapsed
   
    # init variables
    block_num = 0
    ret_num = 0
    bytes_sent = 0
    ret_cnt = max_ret
    result_msg = ''
    successful = False
    last_pkt = False
    init_time = time.time()

    try:
        local_file = open (file_name , "rb" )
    except IOError as message:
        result_msg = str(message)
        print '\n' + result_msg
        return False, result_msg, 0, 0
 
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client_socket.settimeout(1)
    Send_WRQ(client_socket, srv_add, file_name)
    print ('Writting file "{}" to server "{}"'.format(file_name, srv_add[0]))

    while (ret_cnt > 0):
        try:
            data, source = client_socket.recvfrom(1024)
            if block_num == 0 and (srv_add[0] == source[0]): # learn server port
                srv_add = source
            if srv_add == source: # thats my flow
                tftp_pkt = dpkt.tftp.TFTP(data)
                if tftp_pkt.opcode == dpkt.tftp.OP_ACK:
                    if tftp_pkt.block[0] == block_num: # ACK OK
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                        if last_pkt:
                            successful = True
                            result_msg = 'Transfer successful'
                            break
                        else: # send next data block
                            block_num += 1
                            data_block = local_file.read(512)
                            Send_DATA(client_socket, srv_add, block_num, data_block)
                            if len(data_block) < 512:
                                last_pkt = True
                            ret_cnt = max_ret
                            bytes_sent += len(data_block)
                    elif tftp_pkt.block[0] == (block_num - 1):
                        Send_DATA(client_socket, srv_add, block_num, data_block)
                        ret_cnt -= 1
                        ret_num += 1
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                elif tftp_pkt.opcode == dpkt.tftp.OP_ERR:
                    result_msg = tftp_pkt.errmsg
                    break
            else: # source != srv_add , that is not my flow
                result_msg = tftp_err_msg[EBADID]
                Send_ERROR(sckt, source, EBADID)
                continue
        except socket.timeout:
            if block_num == 0:
                Send_WRQ (client_socket, srv_add, file_name)
            else:      
                Send_DATA(client_socket, srv_add, block_num, data_block)
                ret_num += 1
                print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
            ret_cnt -= 1
        except socket.error as message:
            result_msg = str(message)
            Send_ERROR(client_socket, srv_add, EUNDEF)
            break
        except IOError as message:
            result_msg = str(message)
            Send_ERROR(client_socket, srv_add, ENOSPACE)
            break

    client_socket.close()
    local_file.close()
    total_time = time.time()-init_time
    print '\n' + result_msg
    print ('Sent {} bytes in {:>.3f} seconds. Data Rate {:>.0f} B/s'.format(\
        bytes_sent, total_time, bytes_sent/total_time))
    return successful, result_msg, bytes_sent, total_time

def Receive_File(clt_add, file_name, max_ret=3): # Receive file on server

    # clt_add is a tuple ("ip.addr", TID )
    # returns successful condition, result message, received byte number and time elapsed
       
    # init variables
    block_num = 0 # first block number
    bytes_received = 0
    ret_num = 0
    ret_cnt = max_ret
    result_msg = ''
    last_pkt = False
    init_time = time.time()

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client_socket.settimeout(3)

    try:
        local_file = open (file_name , "wb" )
    except IOError as message:
        result_msg = str(message)
        print '\n' + result_msg
        Send_ERROR(client_socket, clt_add, EACCESS)
        client_socket.close()
        return False, result_msg, 0, 0
 
    Send_ACK(client_socket, clt_add, block_num)
    block_num += 1
    print ('Receiving file "{}" from "{}"'.format(file_name, clt_add[0]))

    while (ret_cnt > 0):
        try:
            data, source = client_socket.recvfrom(1024) # recibir datos
            if clt_add == source: # thats my flow
                tftp_pkt = dpkt.tftp.TFTP(data)
                if tftp_pkt.opcode == dpkt.tftp.OP_DATA:
                    if len(tftp_pkt.data) > 512:
                        result_msg = tftp_err_msg[EBADOP]
                        Send_ERROR(client_socket, clt_add, EBADOP)
                        break
                    elif len(tftp_pkt.data) < 512: # last block
                        last_pkt = True
                        result_msg = 'transfer successful'
                    if tftp_pkt.block[0] == block_num: # new_block
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                        local_file.write(tftp_pkt.data)
                        bytes_received += len(tftp_pkt.data)
                        Send_ACK(client_socket, clt_add, tftp_pkt.block[0])
                        block_num += 1
                        if not last_pkt:
                            ret_cnt = max_ret
                        else:
                            ret_cnt = 1 # for dallying
                    elif tftp_pkt.block[0] == (block_num -1): # retransmission
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                        ret_num += 1
                        ret_cnt -= 1
                        Send_ACK(client_socket, clt_add, tftp_pkt.block[0])
                    else:
                        result_msg = tftp_err_msg[EBADOP]
                        Send_ERROR(client_socket, clt_add, EBADOP)
                        break
                elif tftp_pkt.opcode == dpkt.tftp.OP_ERR:
                    result_msg = tftp_pkt.errmsg
                    break
                else:
                    result_msg = tftp_err_msg[EBADOP]
                    Send_ERROR(client_socket, clt_add, EBADOP)
                    break
            else: # source != clt_add , that is not my flow
                result_msg = tftp_err_msg[EBADID]
                Send_ERROR(client_socket, source, EBADID)
                continue
        except socket.timeout:
            if last_pkt:
                break
            else:      
                Send_ACK(client_socket, clt_add, (block_num-1))
                print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                ret_num += 1
            ret_cnt -= 1
        except socket.error as message:
            result_msg = str(message)
            Send_ERROR(client_socket, clt_add, EUNDEF)
            break
        except IOError as message:
            result_msg = str(message)
            Send_ERROR(client_socket, clt_add, ENOSPACE)
            break
     
    client_socket.close()
    local_file.close()
    total_time = time.time()-init_time
    print '\n' + result_msg
    print ('Received {} bytes in {:>.3f} seconds. Data Rate {:>.0f} B/s'.format(\
        bytes_received, total_time, bytes_received/total_time))
    return last_pkt, result_msg, bytes_received, total_time

def Send_File(clt_add, file_name, max_ret=3): # send file from server

    # clt_add is ("ip.addr", TID )
    # TID (a.k.a port)
   
    # init variables
    block_num = 1
    bytes_sent = 0
    ret_num = 0
    ret_cnt = max_ret
    result_msg = ''
    successful = False
    last_pkt = False
    init_time = time.time()

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client_socket.settimeout(3)
 
    try:
        local_file = open (file_name , "rb" )
    except IOError as message:
        result_msg = str(message)
        print '\n' + result_msg
        Send_ERROR(client_socket, clt_add, EACCESS)
        client_socket.close()
        return False, result_msg, 0, 0

    data_block = local_file.read(512)
    Send_DATA(client_socket, clt_add, block_num, data_block)
    print ('Sending file "{}" to "{}"'.format(file_name, clt_add[0]))
 
    if len(data_block) < 512:
        last_pkt = True
    bytes_sent += len(data_block)

    while (ret_cnt > 0):
        try:
            data, source = client_socket.recvfrom(1024) # recibir datos
            if clt_add == source: # thats my flow
                tftp_pkt = dpkt.tftp.TFTP(data)
                if tftp_pkt.opcode == dpkt.tftp.OP_ACK:
                    if tftp_pkt.block[0] == block_num: # ACK OK
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                        if last_pkt:
                            successful = True
                            result_msg = 'Transfer successful'
                            break
                        else: # send next data block
                            block_num += 1
                            data_block = local_file.read(512)
                            Send_DATA(client_socket, clt_add, block_num, data_block)
                            if len(data_block) < 512:
                                last_pkt = True
                            ret_cnt = max_ret
                            bytes_sent += len(data_block)
                    elif tftp_pkt.block[0] == (block_num - 1):
                        ret_num += 1
                        print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
                        Send_DATA(client_socket, clt_add, block_num, data_block)
                        ret_cnt -= 1
                elif tftp_pkt.opcode == dpkt.tftp.OP_ERR:
                    result_msg = tftp_pkt.errmsg
                    break
            else: # source != clt_add , that is not my flow
                result_msg = tftp_err_msg[EBADID]
                Send_ERROR(sckt, source, EBADID)
                continue
        except socket.timeout:
            ret_num += 1
            print ('\rBlock = {} || Retr = {}'.format(block_num, ret_num)),
            Send_DATA(client_socket, clt_add, block_num, data_block)
            ret_cnt -= 1
        except socket.error as message:
            result_msg = str(message)
            Send_ERROR(client_socket, clt_add, EUNDEF)
            break
        except IOError as message:
            result_msg = str(message)
            Send_ERROR(client_socket, clt_add, ENOSPACE)
            break
     
    client_socket.close()
    local_file.close()
    total_time = time.time()-init_time
    print '\n' + result_msg
    print ('Sent {} bytes in {:>.3f} seconds. Data Rate {:>.0f} B/s'.format(\
        bytes_sent, total_time, bytes_sent/total_time))
    return successful, result_msg, bytes_sent, total_time


def main():

    # for debugging
    debug = True

    # processing options
    parser = argparse.ArgumentParser(
      prog = 'MTFTP',
      description = 'A kind of TFTP implementation (client and server) by Guillermo Gomez',
      epilog = ' Hope you enjoy. M0M0 ;)')
 
    parser.add_argument ('operation', choices=['R', 'W', 'S'], \
                         help='Operation to perform: Read / Write / Server')
    parser.add_argument ('-f', dest='file', type=str, \
                         help='File to copy')
    parser.add_argument ('-s', dest='server', type=str,\
                         help='Server address')
    parser.add_argument ('-p', dest='port', type=int, default = 69,\
                         help='Service port')
    '''
    parser.add_argument ('-v', '--verbose', dest='verbose', action="count", \
                         default=0, dest='Server', help='Increase verbosity')
    parser.add_argument ('-d', '--debug', action='store_true', \
                         help='Debug mode. Take care.')
    '''
    options = parser.parse_args()

    if options.port not in range(1,65535):
        print ('Error: Invalid port number')
        exit()

    if options.operation == 'R': # read file
        Read_File ((options.server, options.port),options.file)

    elif options.operation == 'W': # read file
        Write_File ((options.server, options.port),options.file)

    else:
        print('Server mode\nPress Ctrl+C to exit\n')
        try:
            server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            local_host = (socket.gethostname(),options.port)
            server_socket.bind (local_host)
            server_socket.settimeout(3)
        except socket.error as message:
            print message
            exit()
        try:
            while True:
                try:
                    data, address = server_socket.recvfrom(4096)
                    tftp_pkt = dpkt.tftp.TFTP(data)
                    if tftp_pkt.opcode == dpkt.tftp.OP_RRQ:
                        Send_File(address, tftp_pkt.filename)
                        print ''
                    elif tftp_pkt.opcode == dpkt.tftp.OP_WRQ:
                        Receive_File(address, tftp_pkt.filename)
                        print ''
                    else:
                        continue
                except socket.timeout:
                    continue
        except KeyboardInterrupt:
            print 'Work Done'
            server_socket.close()


if __name__ == "__main__":
   main()

No comments:

Post a Comment