#!/usr/bin/perl -w
# Found this on 
# http://xaxxon.slackworks.com/viewcvs/viewcvs.cgi/src/tftp-server/tftp.pl?rev=1.5&content-type=text/vnd.viewcvs-markup
#
# Unknown licence, at the moment. Potentially useful for our provisioning 
# project.
use strict;

sub nonblock ( $ );

use Fcntl;
use IO::Socket;
use IO::Select;
use IO::File; # lets you use filehandles like variables


# Set the port to 69 unless otherwise specified
my $Port = 69;
if ( @ARGV == 1 ) {
	$Port = $ARGV [ 0 ];
} elsif ( @ARGV > 1 ) {
	print "Usage: $0 <port>\n";
	exit ( 0 );
}

# Get the CVSROOT out of the specified file
my $CVS_INFO_FILENAME = "cvs_info";
my $CVSROOT;
$CVSROOT = ReadCvsInfo ( $CVS_INFO_FILENAME );

# Set up the main listening port and set to nonblocking
my $Server = IO::Socket::INET->new ( LocalPort => $Port,
									 Proto     => "udp" ) or
	die "Couldn't bind to port: $Port (Ports <= 1024 require root permission)";
nonblock ( $Server );

# hash for holding clients
my %Clients;

# create a select object and add the main port
my $Select = IO::Select->new ( );
$Select->add ( $Server );

# main loop
while ( 1 ) {

    # select from all handles to see if we've got anything incoming

    foreach my $Readable ( $Select->can_read ( ) ) {

		if ( $Readable == $Server ) {

			# We just got a new connection

			my $Data;
			my $Sender = $Server->recv ( $Data, 1024 );

			# Get IP info on client
			my ( $Port, $IP ) = sockaddr_in ( $Sender );
			my $Hostname = gethostbyaddr ( $IP, AF_INET );

			# make a new socket or just toss this connection in the bit bucket if we can't
			my $OutFd = IO::Socket::INET->new ( PeerAddr => $Hostname,
												PeerPort => $Port,
												Proto    => "udp" ) or
													next;
			print "Got new connection\n";
			nonblock ( $OutFd );

			# add our new socket to the select group
			$Select->add ( $OutFd );
			
			
			$Clients{$OutFd} = {};
			$Clients{$OutFd}->{ip} = $IP;
			$Clients{$OutFd}->{port} = $Port;
			$Clients{$OutFd}->{data} = $Data;

			# Create new entry in Clients hash for this client
			$Clients{$OutFd}->{outfd} = $OutFd;

			my ( $Opcode, $Filename, $Mode ) = UnpackRequest ( $Clients{$OutFd}->{data} );
			
			$Clients{$OutFd}->{opcode} = $Opcode;
			$Clients{$OutFd}->{filename} = $Filename;
			$Clients{$OutFd}->{transfermode} = $Mode;


			nonblock $OutFd;
			$Clients{$OutFd}->{outfd} = $OutFd;

			DealWithClient ( $Clients{$OutFd} );

		} else {

			# Got new data on existing connection
			$Readable->recv ( $Clients{$Readable}->{data}, 1024 );

			DealWithClient ( $Clients{$Readable} );

		}

    }

}


sub CleanupClient 
{

    my ( $Client ) = @_;
	
    # delete the client object from the Clients hash based on the index of its outgoing FD
    delete $Clients{$Client->{outfd}};
	
    # Close their data file
    $Client->{file}->close ( ) if exists $Client->{file} and $Client->{file}->opened ( );

    #  remove their outoing FD from the select list
    $Select->remove ( $Client->{outfd} ) if $Select->exists ( $Client->{outfd} );
    
    # close the outgoing connection
    $Client->{outfd}->close ( );

	# clean up the file
	`rm -rf $Client->{basedir}` if exists $Client->{basedir};
	
}

sub DealWithClient {

    my ( $Client ) = @_;

    $Client->{opcode} = unpack ( "n", $Client->{data} );

    if ( $Client->{opcode} == 1 ) {

		# if we got a read request
		$Client->{action} = "read";
		$Client->{lastack} = 0;
		SendData ( $Client );

    } elsif ( $Client->{opcode} == 2 ) {

		# if we got a write request
		$Client->{action} = "write";
		ReceiveData ( $Client );

    } elsif ( $Client->{opcode} == 3 ) {

		# if we got a data block (512 bytes)
		ReceiveData ( $Client );

    } elsif ( $Client->{opcode} == 4 ) {
		
		# if we got an acknowledgment
		
		# break down the ack
		my $AckPacketNumber = UnpackAck ( $Client->{data} );

		if ( $AckPacketNumber < $Client->{lastack} ) {
			# toss it, it's old
		} elsif ( $AckPacketNumber > $Client->{lastack} + 1 ) {
			# that's an error
			SendError();
			$Client->{disconnect} = 1;
		} else {

			# if we're done sending data
			if ( $Client->{alldataread} ) {
				CleanupClient ( $Client );
			} else {
				$Client->{lastack} = $AckPacketNumber;
				SendData ( $Client );
			}
		}

    } elsif ( $Client->{opcode} == 5 ) {

		# if we got an error
		CleanupClient ( $Client );


    } else {
		#die "Unknown opcode '$Client->{opcode}'";
		CleanupClient ( $Client );
    }

}

sub ReceiveData
{

    my ( $Client ) = @_;
	
    # send back an ack of 0 to say we got the write request
	if ( !exists $Client->{file} ) {
		
		if ( -e $Client->{filename} ) {
			
			$Client->{outfd}->send ( CreateError ( 6, "File already exists" ) );
			CleanupClient ( $Client );
			
		}

		$Client->{file} = IO::File->new ( ">$Client->{filename}" );

		# if we couldn't open the file, send back an error
		if ( !$Client->{file}->opened ( ) ) {
			$Client->{outfd}->send ( CreateError ( 2, "Access Violation" ) );

			CleanupClient ( $Client );

		}
		
		# send a packet 0 ack back
		$Client->{outfd}->send ( CreateAck ( 0 ) );
		$Client->{lastack} = 0;
		
    } else {
		
		
		my ( undef, $PacketNumber, $Data ) = UnpackData ( $Client->{data} );

		# if this is a duplicate of an old packet, then disregard it
		if ( $PacketNumber <= $Client->{lastack} ) {
			return;
		} 

		# else if this packet is too far ahead, then error
		if ( $PacketNumber > $Client->{lastack} + 1 ) {

			$Client->{outfd}->
				send (CreateError ( 0, "Received invalid packet number" ) );

			CleanupClient ( $Client );
			return;

		}

		# otherwise this is the packet we want
		$Client->{file}->print ($Data);
		$Client->{outfd}->send ( CreateAck ( $PacketNumber ) );
		$Client->{lastack} = $PacketNumber;

		if ( length $Data < 512 ) {

			$Client->{file}->close ( );
			CleanupClient ( $Client );

		}
		
    }

}

sub SendData {

    my ( $Client ) = @_;

    if ( !exists $Client->{file} ) {
		
		# check the file out of CVS
		GetFileFromCvs ( $Client );

		$Client->{file} = IO::File->new;
		if ( !$Client->{file}->open ( "$Client->{basedir}/$Client->{filename}" ) ) 
		{

			$Client->{outfd}->send ( CreateError ( 1, "File not found" ) );
			CleanupClient ( $Client );
			return ;
		}


    }

    my $Data;
    my $BytesRead = read $Client->{file}, $Data, 512;

    # store the data in case we have to re-send
    $Client->{lastdata} = $Data; 
    $Client->{outfd}->send ( CreateData ( $Client->{lastack} + 1, $Data ) );

    if ( $BytesRead < 512 ) {

		# close the filehandle since we've read everything
		$Client->{file}->close ( );
		
		# set a flag saying we've sent everything, so the next ack disconnects
		$Client->{alldataread} = 1;
		
    }
    
}








# Create TFTP Data packet
# $Packet is the packet sequence number
# $Buffer is the actual data
sub CreateData
{
    my ( $Packet, $Buffer ) = @_;

    my $BufferLength = length ( $Buffer );
    return pack ( "nna$BufferLength", 3, $Packet, $Buffer );

}

sub CreateError
{

    my ( $ErrorCode, $ErrorString ) = @_;

    return pack ( "nnZ*", 5, $ErrorCode, $ErrorString );

}

# Returns the opcode (3), the block number, and data
sub UnpackData
{

    my ( $Data ) = @_;
    return unpack ( "nna*", $Data ); 

}

sub CreateAck
{
    my ( $Block ) = @_;
    return pack ( "nn", 4, $Block );
}

sub UnpackAck
{

    my ( $Data ) = @_;
    my ( undef, $Packet ) = unpack ( "nn", $Data );

    return $Packet;

}

# Returns the opcode, filename, and mode
sub UnpackRequest
{
    my ( $Data ) = @_;
    return unpack ( "nZ*Z*", $Data );
}

# set $Socket to nonblocking
sub nonblock ( $ )
{
    return;
    my ( $Socket ) = @_;

    my $Flags;

    $Flags = fcntl ( $Socket, F_GETFL, 0 );
    
    fcntl ( $Socket, F_SETFL, $Flags | O_NONBLOCK );

}


sub CleanupCvsFilesinTmp
{

	my ( $Client ) = @_;

	`rm -rf /tmp/$Client->{outfd}`;

}


sub ReadCvsInfo
{

	my ( $Filename ) = @_;

	open CVSINFO, $Filename or 
		die "Couldn't open CVS info file: $Filename";

	my $CvsInfo = <CVSINFO>;

	chomp $CvsInfo;
	
	return $CvsInfo;

}


sub GetFileFromCvs
{

	my ( $User ) = @_;

	my $SockPort = $User->{outfd}->sockport();
	$User->{basedir}="/tmp/CVS-TFTPD-$SockPort";

	# Set the CVS info
	$ENV{CVS_RSH}="ssh";
	$ENV{CVSROOT}=$CVSROOT;

	print "CVSROOT = '$CVSROOT'\n";

	# make a directory to put it in and do the actual checkout.  Note that it
	#   always checks out the file into the entire path.
	print "'mkdir $User->{basedir}'\n";
	print `mkdir $User->{basedir}`;
	chdir ( $User->{basedir} );
	print "Doing checkout\n";
	print `cvs checkout $User->{filename}`;
	print "Done checking out\n";

}

sub DeleteTempCvsFile
{

	my ( $User ) = @_;

	print `rm $User->{basedir}/$User->{filename}`;

}