rc4 cipher for MagPi

rc4.cpp

#include <iostream>
#include <fstream>
#include <cstdlib>
#include <cstring>

class State
{
	unsigned char s[256];
	int i, j;

	void swap(int a, int b);

	public:
	unsigned char getbyte(void);
	State(unsigned char key[], int length );
};

State::State(unsigned char key[], int length)
{
	for(int k=0; k<256; k++)
	{
		s[k]=k;
	}

	j=0;
		
	for(i=0; i<256 ; i++)
	{
		j=(j + s[i] + key[i % length]) % 256;
		swap(i, j);
	}

	i=j=0;
}

void State::swap(int a, int b)
{
	unsigned char temp=s[i];
	s[i]=s[j];
	s[j]=temp;
}

unsigned char State::getbyte(void)
{
	i=(i+1)%256;
	j=(j+s[i])%256;
	swap(i, j);
	int index=(s[i]+s[j])%256;
	return s[index];
}

void parseargs(int argc, char ** argv, std::string & key, std::string & file)
{
	bool readkey = false;
	bool readfile = false;
	bool toomanyargs =false;

	for( int i=1 ; i<argc ; i++ )
	{
		std::string arg = argv[i];
		if(!readkey)
		{
			key=arg;
			readkey=true;
		}
		else if(!readfile)
		{
			file=arg;
			readfile=true;
		}
		else
		{
			toomanyargs=true;
		}
	}

	if(toomanyargs || !readfile || !readkey)
	{
		std::cout << "Usage is: " << argv[0] << " key file" << std::endl;
		exit(EXIT_FAILURE);
	}

	return;
}

int gettextkey(unsigned char data[], std::string key)
{
	if(key.length() > 256)
	{
		std::cout << "ASCII key must be 256 characters or less" <<std::endl;
		exit(EXIT_FAILURE);
	}

	size_t i;

	for(i=0; i<key.length(); i++)
	{
		data[i]=key[i];
	}

	return i;
}

void outfilename(std::string &file)
{
	if(file.find(".rc4", file.length()-4) != std::string::npos) //ie, if file ends with ".rc4"
	{
		file.erase(file.length()-4);
	}
	else
	{
		file.append(".rc4");
	}

}

void openfiles(std::fstream & infile, std::fstream & outfile, std::string & filename)
{
	infile.open(filename.c_str(), std::ios::in | std::ios::binary);
	if(!infile.is_open())
	{
		std::cout << filename << " does not exist" << std::endl;
		exit(EXIT_FAILURE);
	}

	outfilename(filename);

	outfile.open(filename.c_str(), std::ios::in);
	if(outfile.is_open()) //file we are going to write exists!
	{
		std::cout << filename << " already exists, aborting to preserve it" <<std::endl;
		exit(EXIT_FAILURE);
	}

	outfile.close();
	outfile.open(filename.c_str(), std::ios::out | std::ios::binary);
	if(!outfile.is_open())
	{
		std::cout << filename << " could not be opened for writing" << std::endl;
		exit(EXIT_FAILURE);
	}
}

int main(int argc, char **argv)
{
	std::string key, file;
	parseargs(argc, argv, key, file);

	int len=0;
	unsigned char keydata[256];

	len=gettextkey(keydata, key);

	State bytestream (keydata, len);

	std::fstream infile;
	std::fstream outfile;

	openfiles(infile, outfile, file);

	char inbyte;
	char outbyte;
	unsigned char streambyte;

	infile.get(inbyte);

	while(!infile.eof())
	{
		streambyte=bytestream.getbyte();
		outbyte=inbyte ^ streambyte;
		outfile.put(outbyte);
		infile.get(inbyte);
	}

	outfile.close();
	infile.close();

	std::cout << "Encryption finished, output to " << file << std::endl;

	return 0;
}