class Kdbx

# File lib/kdbx.rb, line 16
def initialize(**options)
  @password = @keyfile = nil
  @header, @content = Header.new, String.new
  self.password = options[:password] if options.has_key? :password
  self.keyfile  = options[:keyfile]  if options.has_key? :keyfile
end

# File lib/kdbx.rb, line 7
def self.open(filename, **options)
  new(**options).tap do |kdbx|
    File.open filename, "rb" do |file|
      kdbx.header = Header.load file
      kdbx.decrypt_content file.read
    end
  end
end

# File lib/kdbx/attributes.rb, line 40
def compressionflags
  @header[3].unpack("L").first
end

# File lib/kdbx/attributes.rb, line 44
def compressionflags=(flag)
  @header[3] = [flag].pack("L")
end

# File lib/kdbx/attributes.rb, line 15
def credential
  cred = password || String.new
  return cred if keyfile == nil
  data = IO.read keyfile
  if !data.valid_encoding?
    return cred + sha256(data)
  end
  if data.bytesize == 32
    return cred + data
  end
  if data =~ /\A\h{64}\z/
    data = [data].pack("H*")
    return cred + data
  end
  begin
    xpath = "/KeyFile/Key/Data"
    tnd = REXML::Document.new(data).get_text(xpath)
    cred + Base64.decode64(tnd.to_s)
  rescue REXML::ParseException
    cred + sha256(data)
  end
end

# File lib/kdbx/crypto.rb, line 15
def decrypt_content(data)
  data = decrypt data.to_s
  if data.start_with? streamstartbytes
    size = streamstartbytes.bytesize
    data = data.byteslice size..-1
  else
    fail KeyError, "wrong password or keyfile"
  end
  data = decode data
  data = gunzip data if compressionflags == 1
  data = reverse data if innerrandomstreamid == 2
  @content = data
end

# File lib/kdbx/crypto.rb, line 8
def encrypt_content
  data = @content.to_s
  data = obfuscate data if innerrandomstreamid == 2
  data = gzip data if compressionflags == 1
  encrypt streamstartbytes + encode(data)
end

# File lib/kdbx/attributes.rb, line 64
def encryptioniv
  @header[7]
end

# File lib/kdbx/attributes.rb, line 76
def innerrandomstreamid
  @header[10].unpack("L").first
end

# File lib/kdbx/attributes.rb, line 80
def innerrandomstreamid=(id)
  @header[10] = [id].pack("L")
end

# File lib/kdbx/attributes.rb, line 86
def inspect
  super
end

# File lib/kdbx/attributes.rb, line 11
def keyfile=(str)
  @keyfile = File.absolute_path str
end

# File lib/kdbx/attributes.rb, line 48
def masterseed
  @header[4]
end

# File lib/kdbx/attributes.rb, line 6
def password=(str)
  @password = str == nil ? "" : sha256(str)
end

# File lib/kdbx/attributes.rb, line 68
def protectedstreamkey
  @header[8]
end

# File lib/kdbx.rb, line 23
def save(filename)
  secure_write filename do |file|
    file.write header.dump
    file.write encrypt_content
  end
  true
end

# File lib/kdbx/attributes.rb, line 72
def streamstartbytes
  @header[9]
end

# File lib/kdbx/attributes.rb, line 56
def transformrounds
  @header[6].unpack("Q").first
end

# File lib/kdbx/attributes.rb, line 60
def transformrounds=(num)
  @header[6] = [num].pack("Q")
end

# File lib/kdbx/attributes.rb, line 52
def transformseed
  @header[5]
end

# File lib/kdbx/crypto.rb, line 64
def decode(data)
  io = StringIO.new.binmode
  dt = StringIO.new data
  loop do
    t = dt.readpartial 40
    (hash, size) = t.unpack("x4a32L<")
    break io.string if size == 0
    block = dt.readpartial size
    if sha256(block) != hash
      fail "broken file"
    else
      io.write block
    end
  end
rescue TypeError, EOFError
  fail ParseError, "truncated payload"
end

# File lib/kdbx/crypto.rb, line 48
def decrypt(data)
  cipher = OpenSSL::Cipher.new("AES-256-CBC").decrypt
  cipher.iv, cipher.key = encryptioniv, masterkey
  cipher.update(data) + cipher.final
rescue OpenSSL::Cipher::CipherError
  fail KeyError, "wrong password or keyfile"
end

# File lib/kdbx/crypto.rb, line 56
def encode(data)
  StringIO.new.binmode.tap do |io|
    io.write "\x00" * 4 + sha256(data)
    io.write [data.bytesize].pack("L<")
    io.write data + "\x01" + "\x00" * 39
  end.string
end

# File lib/kdbx/crypto.rb, line 42
def encrypt(data)
  cipher = OpenSSL::Cipher.new("AES-256-CBC").encrypt
  cipher.iv, cipher.key = encryptioniv, masterkey
  cipher.update(data) + cipher.final
end

# File lib/kdbx/crypto.rb, line 89
def gunzip(data)
  StringIO.open data do |io|
    gz = Zlib::GzipReader.new io
    [gz.read, gz.close].first
  end
rescue Zlib::GzipFile::Error => e
  fail ParseError, e.message
end

# File lib/kdbx/crypto.rb, line 82
def gzip(data)
  StringIO.open do |io|
    gz = Zlib::GzipWriter.new io.binmode
    gz.write data; gz.close; io.string
  end
end

# File lib/kdbx/crypto.rb, line 35
def masterkey
  cipher = OpenSSL::Cipher.new("AES-256-ECB").encrypt
  cipher.key, key = transformseed, sha256(credential)
  transformrounds.times { key = cipher.update key }
  sha256(masterseed + sha256(key))
end

# File lib/kdbx/attributes.rb, line 92
def nonce
  "\xE8\x30\x09\x4B\x97\x20\x5D\x2A".b
end

# File lib/kdbx/crypto.rb, line 105
def obfuscate(data)
  xpath = "//Value[@Protected='True']"
  doc, seq = REXML::Document.new(data), sequence
  doc.each_element xpath do |ele|
    t = ele.texts.join.bytes
    t.map! { |b| b ^ seq.next }
    t = Base64.encode64 t.pack("C*")
    ele.text = t.strip
  end
  doc.to_s
rescue REXML::ParseException => e
  fail FormatError, e.message
end

# File lib/kdbx/crypto.rb, line 119
def reverse(data)
  xpath = "//Value[@Protected='True']"
  doc, seq = REXML::Document.new(data), sequence
  doc.each_element xpath do |ele|
    t = Base64.decode64(ele.texts.join).bytes
    ele.text = t.map! { |b| b ^ seq.next }.pack("C*")
  end
  doc.to_s
rescue REXML::ParseException => e
  fail ParseError, e.message
end

# File lib/kdbx.rb, line 33
def secure_write(name)
  name  = File.absolute_path name
  index = -1 - File.extname(name).length
  temp  = 1.step do |i|
    t = name.dup.insert index, ".#{i}"
    break t unless File.exist? t
  end
  begin
    File.open(temp, "wb") { |file| yield file }
    File.rename temp, name
  ensure
    File.delete temp if File.exist? temp
  end
end

# File lib/kdbx/crypto.rb, line 98
def sequence
  cipher = Salsa20.new sha256(protectedstreamkey), nonce
  Enumerator.new do |e|
    loop { cipher.encrypt("\x00" * 64).each_byte { |b| e << b } }
  end
end

# File lib/kdbx/crypto.rb, line 31
def sha256(data)
  OpenSSL::Digest::SHA256.digest data
end