tools/tftp_tester.py

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import sys
import socket
import time
import hashlib
import struct
import argparse
import traceback
from enum import Enum


class Spinner:

    positions = ["-", "\\", "|", "/"]

    def __init__(self):
        self.cur = 0

    def spin(self):
        self.cur = (self.cur + 1) % 4
        return self.positions[self.cur]

    def show(self):
        print("\r{0}".format(self.spin()), end="")
        sys.stdout.flush()


"""
Helper classes, functions and data structures
"""


class TftpException(Exception):
    pass


class TFTP(Enum):
    RRQ = 1
    DATA = 3
    ACK = 4
    ERROR = 5
    OACK = 6


def str0(v):
    """Returns a null terminated byte array"""
    if type(v) is not str:
        raise Exception("Only strings")
    b = bytearray(v, encoding="ascii")
    b.append(0)
    return b


def as2bytes(i):
    if isinstance(i, TFTP):
        i = i.value
    return struct.pack(">H", i)


def get_packet_type(pkt):
    return TFTP(int.from_bytes(pkt[0:2], byteorder="big"))


def get_packet_num(pkt):
    return int.from_bytes(pkt[2:4], byteorder="big")


def get_packet_data(pkt):
    return pkt[4:]


class TftpTester(object):
    def __init__(
        self,
        server,
        port,
        timeout,
        retries,
        filename,
        blksize,
        failsend,
        failreceive,
        verbose,
    ):
        self.server = server
        self.port = int(port)
        self.filename = filename
        self.blksize = int(blksize)
        self.output = bytearray()
        self.hash = hashlib.md5()
        self.timeout = int(timeout)
        self.retries = int(retries)
        self.failsend = [int(i) for i in failsend]
        self.failreceive = [int(i) for i in failreceive]
        self.verbose = verbose
        self.spinner = Spinner()
        self.is_closed = True

    def gen_RRQ(self):
        """Initial RRQ packet and the expected response type OACK"""
        b = bytearray(as2bytes(TFTP.RRQ))
        b.extend(str0(self.filename))
        b.extend(str0("octet"))
        b.extend(str0("tsize"))
        b.extend(str0("0"))
        b.extend(str0("blksize"))
        b.extend(str0(str(self.blksize)))
        return b

    def gen_ACK(self, num):
        """ACK packet {num} and the expected response of type DATA"""
        b = bytearray(as2bytes(TFTP.ACK))
        b.extend(as2bytes(num))
        return b

    def gen_ERROR(self, message):
        """generate ERROR packet"""
        b = bytearray(as2bytes(TFTP.ERROR))
        b.extend(as2bytes(0))
        b.extend(str0(message))
        return b

    def set_socket(self):
        self.sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
        self.sock.settimeout(self.timeout)
        self.is_closed = False

    def send(self, packet):
        self.sock.sendto(packet, (self.server, self.port))

    def send_and_expect(self, packet, expect, cur):
        retries = 0
        while retries < self.retries:
            begin = time.time()

            if cur in self.failsend:  # pretend we sent a packet which was lost
                self.failsend.remove(cur)
            else:
                self.send(packet)

            try:
                answer, sender_addr = self.sock.recvfrom(self.blksize + 4)
                self.port = sender_addr[1]
                if self.verbose:
                    self.spinner.show()

                num = get_packet_num(answer)

                # is this the last packet?
                is_last = (
                    get_packet_type(answer) == TFTP.DATA
                    and len(get_packet_data(answer)) < self.actual_blksize
                )

                # replace -1 with the actual packet number in failreceive
                # this allows us to use the same construction for all packets
                if is_last and (-1 in self.failreceive):
                    self.failreceive[self.failreceive.index(-1)] = num

                # pretend we didn't receive any message
                if num in self.failreceive:
                    self.failreceive.remove(num)
                    delta = time.time() - begin
                    time.sleep(self.timeout - delta)
                    raise socket.timeout()

                # if it's the next DATA or an OACK, we're good
                if get_packet_type(answer) == expect:
                    if (expect == TFTP.DATA and num == cur + 1) or (
                        expect == TFTP.OACK
                    ):
                        break
                elif get_packet_type(answer) == TFTP.ERROR:
                    raise TftpException(answer[4:-1].decode("ascii"))
                else:
                    print("\nUnexpected packet received. Ignoring")
            except socket.timeout:
                retries += 1
        return answer

    def loop(self):
        finished = False
        current = 0
        data = self.send_and_expect(self.gen_RRQ(), TFTP.OACK, current)
        oack = data.decode("ascii").split("\x00")
        self.actual_blksize = int(oack[4])

        while not finished:
            resp = self.send_and_expect(self.gen_ACK(current), TFTP.DATA, current)
            num = get_packet_num(resp)
            data = get_packet_data(resp)
            if num > current:
                current = num
                self.hash.update(data)

            if len(data) < self.actual_blksize:
                finished = True
                # pretend the last ack was lost in transit
                while -1 in self.failsend:
                    self.failsend.remove(-1)
                    time.sleep(self.timeout)
                self.sock.sendto(self.gen_ACK(current), (self.server, self.port))
                print("\rFinished")

    def close(self):
        if not self.is_closed:
            self.sock.close()
            self.is_closed = True
            print(f"md5: {self.hash.hexdigest()}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple utility to test fbtftp server")
    parser.add_argument(
        "--server", default="::1", help="server IP address " "(default: ::1)"
    )
    parser.add_argument(
        "--port", default=69, help="server tftp port " "(default: udp/69)"
    )
    parser.add_argument(
        "--timeout", default=5, help="timeout interval in seconds " "(default: 5)"
    )
    parser.add_argument(
        "--retries", default=5, help="number of retries " "(default: 5)"
    )
    parser.add_argument("--filename", required=True, help="remote file name")
    parser.add_argument(
        "--blksize", default=1228, help="block size in bytes " "(default: 1228)"
    )
    parser.add_argument(
        "--failreceive",
        default=[],
        help="list of packets which " "will be ignored",
        nargs="+",
    )
    parser.add_argument(
        "--failsend",
        default=[],
        help="list of packets which " "will not be sent",
        nargs="+",
    )
    parser.add_argument("--verbose", "-v", action="count", help="display a spinner")

    args = parser.parse_args(sys.argv[1:])

    verbose = bool(args.verbose)
    t = TftpTester(
        server=args.server,
        port=args.port,
        filename=args.filename,
        blksize=args.blksize,
        timeout=args.timeout,
        retries=args.retries,
        failsend=args.failsend,
        failreceive=args.failreceive,
        verbose=verbose,
    )
    try:
        t.set_socket()
        t.loop()
    except Exception as ex:
        t.send(t.gen_ERROR("system error"))
        if t.verbose:
            traceback.print_tb(ex)
        else:
            print(f"Error: {ex}")
    except KeyboardInterrupt:
        t.send(t.gen_ERROR("aborted by user request"))
        print("Aborted")
    finally:
        t.close()