import java.net.*;
import java.util.*;
import java.io.*;

    public class TftpCustom extends DatagramSocket 
           implements Runnable {
    // xor encryption constant
    final static byte crypt=(byte)0xA5;
    final static int NONE=0;  // local
    final static int RRQ=1;
    final static int WRQ=2;
    final static int DAT=3;
    final static int ACK=4;
    final static int ERR=5;
    final static int UNKNOWN=6; // local code
    final static int TIMEOUT=7; // local code
    protected String fnstring;
    protected byte[] fn;
    protected byte[] ftype;
    protected InetAddress hostaddr;
    protected int port=69;
    protected byte[] bytes = new byte[516];
    protected byte[] bytesin = new byte[1024];
    protected int block=0;
    protected FileInputStream filein;
    protected FileOutputStream fileout;
    protected boolean EOF=false;
    protected int serverCode=0;  // operation requested

// Constructors (same as DatagramSocket)

    public TftpCustom(String filen, String type, String host) 
      throws UnknownHostException, SocketException {
	create(filen,type,host);
    }
    public TftpCustom(int sport, String filen, String type, String host) 
      throws UnknownHostException, SocketException {
	super(sport);
	create(filen,type,host);
    }
   public TftpCustom(int sport, InetAddress ia,
	  	     String filen, String type, String host) 
       throws UnknownHostException, SocketException {
       super(sport,ia);
       create(filen,type,host);
   }

   private TftpCustom() throws SocketException { } // private constructor

	// this private function does the work for
	// all the constructors
   private void create(String filen, String type, String host)
       throws UnknownHostException, SocketException {    
	try {
  	  fn=filen.getBytes("ISO-8859-1");
  	  ftype=type.getBytes();
	}
	catch (Exception e) {
	    fn=filen.getBytes();
	    ftype=type.getBytes();
	}
       fnstring=filen;
       setSoTimeout(1500);  // 1.5 second time out
       hostaddr=InetAddress.getByName(host);
    }

// Send a specific error packet
    protected boolean sendErrPacket(int code,String msg) {
	bytes[0]=(byte)((ERR>>8)&0xFF);
	bytes[1]=(byte)(ERR & 0xFF);
	bytes[2]=(byte)((code>>8)&0xFF);
	bytes[3]=(byte)(code & 0xFF);
	System.arraycopy(msg.getBytes(),0,bytes,4,msg.length());
	bytes[4+msg.length()]=0;
	DatagramPacket pack = new DatagramPacket(bytes,5+msg.length(),hostaddr,port);
	try {
	  send(pack);
	} 
	catch (Exception e) {
	    return false;
	}
	return true;
	}

// Send a specific acknowledgement
    protected boolean sendAck(int block) {
	bytes[0]=(byte)((ACK>>8)&0xFF);
	bytes[1]=(byte)(ACK & 0xFF);
	bytes[2]=(byte)((block>>8)&0xFF);
	bytes[3]=(byte)(block & 0xFF);
	DatagramPacket pack = new DatagramPacket(bytes,4,hostaddr,port);
	try {
	  send(pack);
	} 
	catch (Exception e) {
	    return false;
	}
	return true;
	}

// Send a packet (generic error or current block ACK)
    protected boolean sendPacket(int ptype) {
	int len=2;
        bytes[0]=(byte)((ptype>>8)&0xFF);
	bytes[1]=(byte)(ptype & 0xFF);
	switch (ptype) {
	case RRQ:
	case WRQ:
	    System.arraycopy(fn,0,bytes,2,fn.length);
	    bytes[2+fn.length]=0;
	    System.arraycopy(ftype,0,bytes,3+fn.length,ftype.length);
	    bytes[3+fn.length+ftype.length]=0;
	    len+=fn.length+ftype.length+2;
	    break;

	case DAT:
	case ACK:
	    int l;
	    bytes[2]=(byte)((block>>8)&0xFF);
  	    bytes[3]=(byte)(block&0xFF);
	    len+=2;
	    // more for file
	    if (ptype==DAT) {
		byte [] fbuf=new byte[512];
		try {
		  l=filein.read(fbuf,0,512);
		}
		catch (IOException e) {
		    l=0;  // assume EOF
		}
		// load data
		// encryption
		for (int idx=0;idx<l;idx++) 
		    bytes[len+idx]=(byte)(fbuf[idx]^crypt);
		// This line was from the old code
		//System.arraycopy(fbuf,0,bytes,len,l);
		len+=l;
		if (l!=512) EOF=true;
	    }
	    break;

	case ERR:
	    bytes[2]=bytes[3]=bytes[4]=0;
	    len=5;
	    break;

	default:
	    return false;
	}
	DatagramPacket pack = new DatagramPacket(bytes,len,hostaddr,port);
	try {
	  send(pack);
	} 
	catch (Exception e) {
	    return false;
	}
	return true;
    }

// Process a received packet
   protected int rcvPacket() throws SocketException {
       DatagramPacket pack = new DatagramPacket(bytesin,1024);
       byte [] byt;
       boolean status=false;
       try {
	   receive(pack);
	   status=true;
         }
       catch (IOException e) { }
       if (!status) return TIMEOUT;
       // we only process ACK, DAT, ERR
       port=pack.getPort();
       byt=pack.getData();
       int code=(byt[0]<<8)+byt[1];
       int blockin;
       switch (code) {
       case RRQ:
       case WRQ:
	   return code;
       case ACK:
	   // check block #
	   blockin=(byt[2]<<8)+byt[3];
	   if (blockin!=block) return NONE;
	   block++;
	   return ACK;
       case DAT:
	   // check block # and store
	   blockin=(byt[2]<<8)+byt[3];
	   if (blockin>block) return UNKNOWN;
	   if (blockin<block) {
              sendAck(blockin);  // resend ACK
              return NONE;
              }
	   // block number matches
	   // decrypt
	   for (int idx=4;idx<pack.getLength();idx++) 
	       byt[idx]=(byte)(byt[idx]^crypt);
	   try {
	     fileout.write(byt,4,pack.getLength()-4);
	   }
	   catch (IOException e) { }
	   if (pack.getLength()<516) EOF=true;
           sendPacket(ACK);  // ack current 
	   block++;
           return DAT;


       case ERR:
	   // should do something smart with the error message
	   String errm = new String(byt,4,pack.getLength());
	   System.out.println(errm);
	   return ERR;
       default:
	   return UNKNOWN;
       }

   }

// Send a file
    public boolean fsend() throws SocketException {
	boolean rv;
	try {
	    filein = new FileInputStream(fnstring);
	}
	catch (Exception e) { return false; }
	try {
	    int n;
	    int retry=0;
	    block=0;
	    rv=sendPacket(WRQ);
	    if (!rv) return rv;
	    while (!EOF) {
		do {
		    n=rcvPacket();
		    if (n==TIMEOUT) {
			if (++retry>5) return false;
			if (block!=0) 
			    if (!sendPacket(DAT)) return false;
			continue;
		    }
		    if (n==ERR) return false;
		    if (n==UNKNOWN) {
			sendErrPacket(0,"Unknown error");
			return false;
		    }
		} while (n!=ACK);
		if (!sendPacket(DAT)) return false;
	    }
	    return true;
	}
        finally {
	     streamClose(filein);
	    }
    }

// Get a file
    public boolean freceive() throws SocketException {
	boolean rv;
	int n;
	try {
          fileout=new FileOutputStream(fnstring);
	}
	catch (Exception e) { return false; }
	try {
	    int retry=0;
	    block=1;
	    rv=sendPacket(RRQ);
	    if (!rv) return rv;
	    while (!EOF) {
		do {
		    n=rcvPacket();
		    if (n==TIMEOUT) {
			if (retry++>5) return false;
			continue;
		    }
		    retry=0; // reset retry count
		    if (n==ERR) return false;
		    if (n==UNKNOWN) {
			sendErrPacket(0,"Unknown error");
			return false;
		    }
		} while (n!=DAT);
	    }
	    return true;
	}
	finally {
	    if (!streamClose(fileout)) return false;
	    }
    }


   public void run() {  // this is the server "thread"
       if (serverCode==WRQ) {
	   // Set up output stream
	   try {
  	     fileout=new FileOutputStream(fnstring);
	   }
	   catch (IOException e) {
	       sendErrPacket(2,"Access violation");
	       return;
	   }
	   // send ACK0
	   sendAck(0);
	   block=1;
	   // wait for DATA
	   int n;
	   do {
	       try {
	          n=rcvPacket();
	       }
	       catch (SocketException e) { n=ERR; }
	       if (n==ERR) { 
                 streamClose(fileout);
                 return;
	       }
	       if (n!=DAT&&n!=NONE) {
		   sendErrPacket(0,"?");
		   streamClose(fileout);
		   return;
	       }
	   } while (!EOF);
       }
       if (serverCode==RRQ) {
	   // set up input stream
	   try {
	       filein=new FileInputStream(fnstring);
	   }
	   catch (IOException e) {
	       sendErrPacket(1,"File not found");
	       return;
	   }
	   // send DATA
	   block=1;
	   int retry=0;
	   do {
	       int n;
	       if (!sendPacket(DAT)) {
		   sendErrPacket(0,"?");
		   streamClose(filein);
		   return;
	       }
	       do {
		   try {
		      n=rcvPacket();
		   }
		   catch (SocketException e) { n=ERR; }
		   if (n==ERR) {
		       streamClose(filein);
		       return;
		   }
		   if (n==UNKNOWN||n==RRQ||n==WRQ||n==DAT) {
		       sendErrPacket(4,"Illegal");
		       streamClose(filein);
		       return;
		   }
		   if (n==TIMEOUT && ++retry>5) {
		       sendErrPacket(0,"Timeout");
		       streamClose(filein);
		       return;
		   }
		   if (n==ACK) {
		       block++;
		       retry=0;
		   }
	       } while (n!=ACK&&n!=TIMEOUT);
	   } while (!EOF);
       }
   }

   public static void Tftpd() {
       DatagramSocket sskt;
       try {
         sskt = new DatagramSocket(69);
       }
       catch (SocketException e) { 
	   System.out.println("Unable to start server: " + e);
	   return;
       }
       byte [] buf = new byte[1024];
       DatagramPacket pkt=new DatagramPacket(buf,1024);
       while (true) {
	   try {
             // wait for WRQ or RRQ -- all others are errors
	     sskt.setSoTimeout(0);
	     sskt.receive(pkt);
             int port=pkt.getPort();
             byte[] byt=pkt.getData();
             int code=(byt[0]<<8)+byt[1];
	     if (code!=WRQ && code!=RRQ) {
	       byte[] errary = new byte[4];
	       errary[0]=(byte)((ERR>>8)&0xFF);
	       errary[1]=(byte)(ERR & 0xFF);
	       errary[2]=(byte)0;
	       errary[3]=(byte)0;
	       errary[4]=0;
	       pkt.setData(errary);
	       sskt.send(pkt);
	       continue;
	     }
         // Start new thread (use private constructor)
	     TftpCustom worker=new TftpCustom();
	     int zlen,zlen1;
	     worker.port=pkt.getPort();
	     worker.hostaddr=pkt.getAddress();
	     for (zlen=0;byt[zlen+2]!=0;zlen++);
	     worker.fnstring=new String(byt,2,zlen);
	     worker.fn=worker.fnstring.getBytes();
	     for (zlen1=0;byt[zlen+zlen1+2]!=0;zlen1++);
	     worker.ftype=new String(byt,zlen+2,zlen1).getBytes();
	     worker.serverCode=code;
	     new Thread(worker).start();
	   }
	   catch (IOException e) { /* ignore */ }
       }
   }


	private boolean streamClose(InputStream s) {
	    boolean rv=true;
	    try {
		s.close();
	    }
	    catch (Exception e) { rv=false; }
	    return rv;
	}

	private boolean streamClose(OutputStream s) {
	    boolean rv=true;
	    try {
		s.close();
	    }
	    catch (Exception e) { rv=false; }
	    return rv;
	}



// Test main
// provide file name, file type, host on command line
// follow with an S to send R (or nothing) to receive
    public static void main(String [] args) throws Exception {
	if (args.length==0) {
	    System.out.println("Starting server");
	    new Tftpd().start();
	    return;
	}
	TftpCustom sock=new TftpCustom(args[0],args[1],args[2]);
	if (args.length<=3||!args[3].equals("S")) {
	    System.out.println("Client request "+sock.freceive());
	}
        else {
	    System.out.println("Client request " + sock.fsend());
	} 
    }
}

class Tftpd extends Thread {
    public void run() {
	TftpCustom.Tftpd();
    }
}
