require 'optparse'

class Integer
	def letter?
		self <= ?Z and self >= ?A
	end
end

class String
	def letters
		delete '^A-Z'
	end
	
	def map_letter
		unpack('C*').inject('') {|s, b| s << (b.letter? ? yield(b) : b)}
	end

	def substr(len, phase)
		scan(/#{'.'*len}/).inject('') { |ss, c| ss << c[phase] }
	end
	
	def decipher(key)
		k = key.dup
		idx = -1
		map_letter { |l|
			idx = (idx + 1) % key.length
			k[idx] = (l - k[idx]) % 26 + ?A
		}
	end
	
	def IC
		count = Hash.new(0.0)
		letters.each_byte { |l| count[l] += 1 }
	
		tot = count.values.inject { |a, b| a + b }
		(?A .. ?Z).map{ |l| count[l] }.map{ |l| l * (l-1) }.inject { |a, b| a + b } / (tot * (tot-1))
	end
end

def padv s
	print s + "\r"
	$stdout.flush
	print " "*s.length + "\r"
end

def find_keys_ic(str, threshold, keyleninterval = 7..12, verbose = true)
	keyleninterval.map { |keylen|
		puts "Trying keylen #{keylen}" if verbose
		h = {}
		('A'..'Z').each { |key|
			ic = str.substr(keylen, 0).decipher(key).IC
			h[key] = ic if ic > threshold
		}
		next if h.empty?

		puts "[+] Found something !" if verbose
		
		keychars = []
		keylen.times { |kl|
			keychars << []
			('A'..'Z').each { |key|
				padv((keychars + [[key]]).map { |k| k.join }.join(', '))
				keychars.last << key if str.substr(keylen, kl).decipher(key).IC > threshold
			}
		}

		puts "[+] Key chars: #{keychars.map { |k| k.join }.join(', ')}" if verbose

		keychars.inject(['']) { |all, cur|
			cur.map { |chr|
				padv all.join(', ') if verbose
				all.map { |a| a + chr }
			}.flatten.select { |pkey|
				# interleave substrings
				ss = (0...pkey.length).map { |i|
					str.substr(keylen, i).unpack('C*')
				}
				ss.shift.zip(*ss).flatten.pack('C*').decipher(pkey).IC > threshold
			}
		}
	}.compact.flatten
end

threshold = 0.07

opts = OptionParser.new
opts.banner = "Usage: #$0 [options] [file]"
opts.on('-k', '--key KEY', 'Decodes the file with KEY') { |k|
	puts ARGF.read.decipher(k)
	exit
}
opts.on('-t', '--threshold T', Float, "Sets the Index of coincidence threshold (defaults to #{threshold})") { |threshold| }

opts.on('-h', '--help', 'Shows this message') {
	puts opts
	exit
}
opts.parse! ARGV

raw = ARGF.read
keys = find_keys_ic(raw.letters, threshold)
#puts "keys with good IC: #{keys.inspect}"
puts
puts "The key must be " + keys.sort_by { |k| raw.decipher(k).count('AEIOU') }.last.inspect