#!/usr/bin/ruby

require 'openssl'
require 'base64'
require 'securerandom'

require 'net/http'
require 'net/https'

def pbkdf2(password, salt, keylen, opts = {})
	hash = opts[:hash] || 'sha1'
	iterations = (opts[:iterations] || 1) - 1
	def bigendian(val, len)
		len.times.map { val, r = val.divmod 256; r }.reverse.pack('C*')
	end
	key = ''
	blockindex = 1
	while key.length < keylen
		block = OpenSSL::HMAC.digest(hash, password, salt + bigendian(blockindex, 4))
		block = block.unpack('C*')
		u = block
		iterations.times do
			u = OpenSSL::HMAC.digest(hash, password, u.pack('C*')).unpack('C*')
			block.length.times.each do |j|
				block[j] ^= u[j]
			end
		end
		block = block.pack('C*')
		key += block
		blockindex += 1
	end
	key.slice(0, keylen)
end

def aes_decrypt(text, key, opts = {})
	text = text.dup
	iv = text.slice!(0, 16) # aes has fixed blocksize of 128 bits
	key = pbkdf2(key, iv, (opts[:aeskeysize] || 256) / 8, opts)

	cipher = OpenSSL::Cipher::AES.new(8 * key.length, opts[:mode] || :OFB).decrypt
	cipher.key = key
	cipher.iv = iv
	cipher.update(text) + cipher.final
end

def aes_encrypt(text, key, opts = {})
	cipher = OpenSSL::Cipher::AES.new(opts[:aeskeysize] || 256, opts[:mode] || :OFB).encrypt
	cipher.iv = iv = opts[:iv] || cipher.random_iv
	cipher.key = pbkdf2(key, iv, (opts[:aeskeysize] || 256) / 8, opts)
	text = cipher.update(text) + cipher.final
	iv + text
end

def aes_256_ofb_decrypt(text, key, opts = {})
	opts = opts.merge :mode => :OFB, :aeskeysize => 256
	aes_decrypt(text, key, opts)
end

def aes_256_ofb_encrypt(text, key, opts = {})
	opts = opts.merge :mode => :OFB, :aeskeysize => 256
	aes_encrypt(text, key, opts)
end


def aes_128_cbc_key_iv_derive(key, salt)
	t = key + salt
	aeskey = OpenSSL::Digest::MD5::digest t
	iv = OpenSSL::Digest::MD5::digest (aeskey + t)
	[aeskey, iv]
end

def aes_128_cbc_decrypt(text, key, opts = {})
	text = text.dup
	salt = text.slice!(0, 16) # 8 bytes 'Salted__', 8 bytes salt
	magic = salt.slice!(0, 8)
	throw 'invalid input, missing salt' if (magic != 'Salted__')
	cipher = OpenSSL::Cipher::AES.new(128, :CBC).decrypt
	cipher.key, cipher.iv = aes_128_cbc_key_iv_derive(key, salt)
	(cipher.update(text) + cipher.final)[0..-2] # drop trailing \n
end

def aes_128_cbc_encrypt(text, key, opts = {})
	cipher = OpenSSL::Cipher::AES.new(128, :CBC).encrypt
	salt = cipher.random_iv.slice(0, 8) # only 8 from 16 bytes
	cipher.key, cipher.iv = aes_128_cbc_key_iv_derive(key, salt)
	'Salted__' + salt + cipher.update(text + "\n") + cipher.final
end


def ccm_start_iv(nonce, taglength)
	q = 15 - nonce.length
	[q-1, nonce, [0].pack('C') * q].pack('Ca*a*')
end

def ccm_tag(cipherMac, text, nonce, taglength)
	cipherMac.reset
	q = 15 - nonce.length
	msglen = [ text.length ].pack('Q>')[-q..-1]
	b0 = [ 4*(taglength - 2) + (q-1), nonce, msglen ].pack('Ca*a*')
	cipherMac.update(b0 + text + [0].pack('C') * 15).slice(-16, taglength)
end

def ccm_decrypt(cipher, cipherMac, text, nonce, taglength)
	cipher.reset
	cipher.iv = ccm_start_iv(nonce, taglength)
	# need to decrypt tag first
	tag = text.slice!(-taglength, taglength)
	text = cipher.update(tag + text) + cipher.final
	tag = text.slice!(0, taglength)
	raise "Invalid tag" if tag != ccm_tag(cipherMac, text, nonce, taglength)
	text
end

def ccm_encrypt(cipher, cipherMac, text, nonce, taglength)
	cipher.reset
	cipher.iv = ccm_start_iv(nonce, taglength)
	# need to encrypt tag first
	tag = ccm_tag(cipherMac, text, nonce, taglength)
	text = cipher.update(tag + text) + cipher.final
	# put tag at the end
	tag = text.slice!(0, taglength)
	text + tag
end

def aes_256_ccm_key_nonce_derive(key, salt)
	aeskey = pbkdf2(key, salt, 32 + 8, { :hash => 'sha256', :iterations => 100 })
	nonce = aeskey.slice!(32, 8)
	[aeskey, nonce]
end

def aes_256_ccm_decrypt(text, key, opts = {})
	text = text.dup
	salt = text.slice!(0, 8)
	aeskey, nonce = aes_256_ccm_key_nonce_derive(key, salt)
	cipher = OpenSSL::Cipher::AES.new(256, :CTR).encrypt
	cipher.key = aeskey
	cipherMac = OpenSSL::Cipher::AES.new(256, :CBC).encrypt
	cipherMac.key = aeskey
	ccm_decrypt(cipher, cipherMac, text, nonce, 16)
end

def aes_256_ccm_encrypt(text, key, opts = {})
	cipher = OpenSSL::Cipher::AES.new(256, :CTR).encrypt
	cipherMac = OpenSSL::Cipher::AES.new(256, :CBC).encrypt
	salt = opts[:salt] || cipher.random_iv.slice(0, 8)
	aeskey, nonce = aes_256_ccm_key_nonce_derive(key, salt)
	cipher.key = aeskey
	cipherMac.key = aeskey
	salt + ccm_encrypt(cipher, cipherMac, text, nonce, 16)
end


def generateKey(len = 24)
	chars = 'abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
	len.times.map { chars[SecureRandom.random_number(chars.length)].ord }.pack('c*')
end

def decrypt(text, key, cipher, opts = {})
	text = Base64.decode64(text)
	case cipher
	when 'AES-256-OFB'
		aes_256_ofb_decrypt(text, key, opts)
	when 'AES-128-CBC'
		aes_128_cbc_decrypt(text, key, opts)
	when 'AES-256-CCM'
		aes_256_ccm_decrypt(text, key, opts)
	else
		throw ('Unknown cipher "' + cipher + '"')
	end
end

def encrypt(text, cipher, opts = {})
	key = opts[:key] || generateKey(opts[:keylen] || 24)
	ctext = case cipher
	when 'AES-256-OFB'
		aes_256_ofb_encrypt(text, key, opts)
	when 'AES-128-CBC'
		aes_128_cbc_encrypt(text, key, opts)
	when 'AES-256-CCM'
		aes_256_ccm_encrypt(text, key, opts)
	else
		throw ('Unknown cipher "' + cipher + '"')
	end
	ctext = Base64.encode64(ctext).gsub(/\s+/, "")
	raise "decryption failed" if decrypt(ctext, key, cipher, opts) != text
	[ ctext, key ]
end

def fixipv6host(host)
	if m = /^\[([0-9a-fA-F:]+)\]$/.match(host)
		return m[1]
	end
	host
end

def http_get(uri)
	request = Net::HTTP::Get.new uri.request_uri
	request['Host'] = uri.host
	http = Net::HTTP.new(fixipv6host(uri.host), uri.port)
	http.use_ssl = uri.scheme == 'https'
# 	http.verify_mode = OpenSSL::SSL::VERIFY_PEER
# 	http.verify_depth = 5
	http.verify_mode = OpenSSL::SSL::VERIFY_NONE
	http.start { |http| http.request request }
end

def http_post(uri, args)
	request = Net::HTTP::Post.new(uri.request_uri)
	request['Host'] = uri.host
	request.set_form_data(args)
	http = Net::HTTP.new(fixipv6host(uri.host), uri.port)
	http.use_ssl = uri.scheme == 'https'
# 	http.verify_mode = OpenSSL::SSL::VERIFY_PEER
# 	http.verify_depth = 5
	http.verify_mode = OpenSSL::SSL::VERIFY_NONE
	http.start { |http| http.request request }
end

def hashpw(password)
	return nil if password.nil?
	return OpenSSL::Digest::SHA1.hexdigest(password)
end

require 'optparse'
opts = {}
OptionParser.new do |o|
	o.on('-u', '--url URL', "Retrieve paste from url (conflicts with the posting options)") { |url| opts[:url] = URI(url) }
	o.on('-f', '--file FILENAME', "Upload file") { |fn| opts[:fn] = fn }
	o.on('-m', '--mime MIMETYPE', "Specify mime type for paste") { |mime| opts[:mime] = mime }
	o.on('-t', '--ttl TTL', "Specify Time-To-Live for paste in seconds, default one week (-1 for indefinately)") { |ttl| opts[:ttl] = ttl }
	o.on('-p', '--password[PASSWORD]', "Use password protection on server side (no additional encryption)") { |password|
		opts[:pass] = true
		opts[:password] = password unless password.nil?
	}
	o.on('-s', '--site SITE', "Post upload to another ezcrypt pastebin (default: https://ezcrypt.it)") { |s| opts[:site] = s }
	o.separator ""
	o.separator "    If neither url nor filename was given, a final parameter can be used to specify it. Urls are autodetected."
	o.separator ""
	o.on('-h', '--help', "Show this help") { STDERR.puts o; exit }
	o.parse!
	if ARGV.length == 1
		if opts[:url] || opts[:fn]
			STDERR.puts o
			exit 1
		end
		begin
			url = URI(ARGV[0])
			if ("http" == url.scheme || "https" == url.scheme) && url.host
				opts[:url] = url
			else
				opts[:fn] = ARGV[0]
			end
		rescue
			opts[:fn] = ARGV[0]
		end
	end
	if ARGV.length > 1 or (opts[:url] and (opts[:fn] || opts[:ttl] || opts[:mime] || opts[:site] || opts[:pass])) or (!opts[:url] and !opts[:fn])
		STDERR.puts o
		exit 1
	end
end

if opts[:pass] and opts[:password].nil?
	if opts[:fn] == '-'
		STDERR.puts "Can't read post data from stdin and prompt for password"
		exit 1
	end
	STDERR.write "Enter password: "
	STDERR.flush
	opts[:password] = STDIN.readline.chomp
end

if opts[:url]
	uri = opts[:url]
	if !uri.fragment
		STDERR.puts "Specified url has no fragment, cannot decode paste"
		exit 1
	end

	uri.query = (uri.query || '') + '&raw'
	key = uri.fragment
	uri.fragment = nil

	password = nil
	while
		if password.nil?
			resp = http_get(uri)
		else
			resp = http_get(uri, :p => hashpw(password))
		end

		if 403 == resp.code.to_i
			STDERR.write "Paste is password protected. Enter password: "
			STDERR.flush
			password = STDIN.readline.chomp
		elsif 200 != resp.code.to_i
			STDERR.puts "Got HTTP/#{resp.http_version} #{resp.code} #{resp.message}"
			STDERR.puts "Location: #{resp['Location']}" if resp['Location']
			exit 1
		end

		# STDERR.puts ("Syntax: " + resp.header['X-Syntax'].to_s)
		cipher = resp.header['X-Cipher'] || 'AES-256-OFB'
		STDOUT.write decrypt(resp.body, key, cipher)
		exit 0
	end

else
	uri = URI(opts[:site] || 'https://ezcrypt.it')
	if opts[:fn] != '-'
		if opts[:mime].nil?
			begin
				opts[:mime] = IO.popen("file --brief --mime-type '#{opts[:fn]}'", "r").read.chomp
			rescue
				# use default mime type text/plain
			end
		end
		text = File.open(opts[:fn], "rb").read
	else
		text = STDIN.read
	end

	cipher = 'AES-256-CCM'
	ciphertext, key = encrypt(text, cipher)
	resp = http_post(uri, :data => ciphertext, :cipher => cipher, :syn => opts[:mime] || 'text/plain', :ttl => opts[:ttl] || (7*86400), :p => hashpw(opts[:password]))
	if 200 != resp.code.to_i
		STDERR.puts "Key is #{key}"
		STDERR.puts "Got HTTP/#{resp.http_version} #{resp.code} #{resp.message}"
		resp.each_header { |k,v| STDERR.puts "#{k}: #{v}" }
		STDERR.puts resp.body
		exit 1
	end
	html = resp.body
	if html =~ /^\{"id":".*"\}\s*$/m then
		id = html.gsub(/^\{"id":"/m, "").gsub(/"\}\s*$/m, "")
		uri.path += '/' unless ?/ == uri.path[-1]
		uri.path += id
		uri.fragment = key
		puts uri
	else
		STDERR.puts "Key is #{key}"
		STDERR.puts "Can't parse response #{resp.body}"
		exit 1
	end
end