Added checks to ensure the input file contains enough unique words to…

… generate

the specified phrases file

smpp_generator.py

@@ -21,43 +21,53 @@
 """
 
 from random import choice
+from itertools import permutations
+from math import factorial
 import argparse
 
+
+
 def build_wordlist(filename):
-
-	words = [] 
-
-	wordfile = open(filename,'r')
-
-	for line in wordfile:
-		for word in line.split():
-			if words.count(word) is 0:
-				words.append(word) 	
-
-	wordfile.close()
-	return words
+
+	try:	
+		words = [] 
+		wordfile = open(filename,'r')
 
 
+		for line in wordfile:
+			for word in line.split():
+				if words.count(word) is 0:
+					words.append(word) 	
+
+		wordfile.close()
+		return words
+
+	except IOError,e:
+		print "filename {:s} does not exist".format(filename)
+		exit()
+
 def generate_phrase(wordlist,numofwords):
-
-	phrase = []		
+	try:	
+		phrase = []		
 
-	for i in range(0,numofwords):
-		word = choice(wordlist)	
-
-		#Ensure one occurrence of each word	
-		while phrase.count(word) is not 0:
+		for i in range(0,numofwords):
 			word = choice(wordlist)	
+
+			#Ensure one occurrence of each word	
+			while phrase.count(word) is not 0:
+				word = choice(wordlist)	
 
-		phrase.append(word)
-
-	pphrase = ""
-
-	for i in range(0,len(phrase)):
-		pphrase = pphrase +" "+ phrase[i]		
-
-	return pphrase
-
+			phrase.append(word)
+
+		pphrase = ""
+
+		for i in range(0,len(phrase)):
+			pphrase = pphrase +" "+ phrase[i]		
+
+		return pphrase
+	except IndexError,e:
+		print 'File is empty'
+		exit()
 
 def write_phrase_to_file(phrase,filename):
 	phrasefile = open(filename,'a')
@@ -67,7 +77,19 @@ def write_phrase_to_file(phrase,filename):
 def	generate_phrase_file(noOfWords,noOfLines,inFile,outFile):
 	phraselist =[]	
 	wordlist = build_wordlist(inFile)	
+
+	numlist = range(0,(len(wordlist)))
+	count = 0
 
+	for p in permutations(numlist,noOfWords):
+		count+=1
+		if count == noOfLines:
+			break
+
+	if count < noOfLines:
+		print "Not enough unique words in input file {:s}".format(inFile)
+		exit()
+
 	zeros = len(str(noOfLines));
 
 
@@ -91,12 +113,20 @@ def main():
 			" comes to verify during an OTR or similar session, the question can"
 			" simply be the number and the answer is the phrase")
 
-	parser.add_argument("noOfWords",type=int,help="The number of words desired in the phrase(s)")
-	parser.add_argument("noOfLines",type=int,help="The number of lines of phrases to generate")
+	parser.add_argument("noOfWords",type=int,help="The number of words desired in the phrase(s). Max is 20")
+	parser.add_argument("noOfLines",type=int,help="The number of phrases to generate. Max is 500")
 	parser.add_argument("inFile",type=str,help="The input file used to generate the phrases")
 	parser.add_argument("outFile",type=str,help="The output file to save the phrases to")
 
 	args = parser.parse_args()
+
+	if args.noOfWords > 20:
+		print "The number of words per passphrase cannot be more than 20"
+		exit()
+
+	if args.noOfLines > 500:
+		print "The number of passphrases cannot be more than 500"
+		exit()
 
 	generate_phrase_file(args.noOfWords,args.noOfLines,args.inFile,args.outFile)