Set master key after reading the database.

5 files changed, 68 insertions, 27 deletions

diff --git a/src/format/KeePass1Reader.cpp b/src/format/KeePass1Reader.cpp
index d0ff5fbf4..21448ef81 100644
--- a/src/format/KeePass1Reader.cpp
+++ b/src/format/KeePass1Reader.cpp
@@ -30,6 +30,8 @@
 #include "crypto/CryptoHash.h"
 #include "format/KeePass1.h"
 #include "keys/CompositeKey.h"
+#include "keys/FileKey.h"
+#include "keys/PasswordKey.h"
 #include "streams/SymmetricCipherStream.h"
 
 class KeePass1Key : public CompositeKey
@@ -52,8 +54,26 @@ KeePass1Reader::KeePass1Reader()
 }
 
 Database* KeePass1Reader::readDatabase(QIODevice* device, const QString& password,
-                                       const QByteArray& keyfileData)
+                                       QIODevice* keyfileDevice)
 {
+    QByteArray keyfileData;
+    FileKey newFileKey;
+
+    if (keyfileDevice) {
+        keyfileData = readKeyfile(keyfileDevice);
+
+        if (keyfileData.isEmpty()) {
+            return 0;
+        }
+        if (!keyfileDevice->seek(0)) {
+            return 0;
+        }
+
+        if (!newFileKey.load(keyfileDevice)) {
+            return 0;
+        }
+    }
+
     QScopedPointer<Database> db(new Database());
     QScopedPointer<Group> tmpParent(new Group());
     m_db = db.data();
@@ -197,22 +217,41 @@ Database* KeePass1Reader::readDatabase(QIODevice* device, const QString& passwor
         entry->setUpdateTimeinfo(true);
     }
 
+    CompositeKey key;
+    if (!password.isEmpty()) {
+        key.addKey(PasswordKey(password));
+    }
+    if (keyfileDevice) {
+        key.addKey(newFileKey);
+    }
+
+    db->setKey(key);
+
     return db.take();
 }
 
 Database* KeePass1Reader::readDatabase(const QString& filename, const QString& password,
-                                       const QByteArray& keyfileData)
+                                       const QString& keyfileName)
 {
-    QFile file(filename);
-    if (!file.open(QFile::ReadOnly)) {
-        raiseError(file.errorString());
+    QFile dbFile(filename);
+    if (!dbFile.open(QFile::ReadOnly)) {
+        raiseError(dbFile.errorString());
         return 0;
     }
 
-    QScopedPointer<Database> db(readDatabase(&file, password, keyfileData));
+    QScopedPointer<QFile> keyFile;
+    if (!keyfileName.isEmpty()) {
+        keyFile.reset(new QFile(keyfileName));
+        if (!keyFile->open(QFile::ReadOnly)) {
+            raiseError(keyFile->errorString());
+            return 0;
+        }
+    }
+
+    QScopedPointer<Database> db(readDatabase(&dbFile, password, keyFile.data()));
 
-    if (file.error() != QFile::NoError) {
-        raiseError(file.errorString());
+    if (dbFile.error() != QFile::NoError) {
+        raiseError(dbFile.errorString());
         return 0;
     }
 
diff --git a/src/format/KeePass1Reader.h b/src/format/KeePass1Reader.h
index 040343681..c90328288 100644
--- a/src/format/KeePass1Reader.h
+++ b/src/format/KeePass1Reader.h
@@ -35,10 +35,9 @@ class KeePass1Reader
 public:
     KeePass1Reader();
     Database* readDatabase(QIODevice* device, const QString& password,
-                           const QByteArray& keyfileData);
+                           QIODevice* keyfileDevice);
     Database* readDatabase(const QString& filename, const QString& password,
-                           const QByteArray& keyfileData);
-    static QByteArray readKeyfile(QIODevice* device);
+                           const QString& keyfileName);
     bool hasError();
     QString errorString();
 
@@ -61,6 +60,7 @@ private:
     bool parseGroupTreeState(const QByteArray& data);
     bool parseCustomIcons4(const QByteArray& data);
     void raiseError(const QString& str);
+    static QByteArray readKeyfile(QIODevice* device);
     static QDateTime dateFromPackedStruct(const QByteArray& data);
     static bool isMetaStream(const Entry* entry);
 
diff --git a/tests/TestKeePass1Reader.cpp b/tests/TestKeePass1Reader.cpp
index dbf11144d..587204131 100644
--- a/tests/TestKeePass1Reader.cpp
+++ b/tests/TestKeePass1Reader.cpp
@@ -17,7 +17,7 @@
 
 #include "TestKeePass1Reader.h"
 
-#include <QtCore/QFile>
+#include <QtCore/QBuffer>
 #include <QtTest/QTest>
 
 #include "config-keepassx-tests.h"
@@ -28,6 +28,10 @@
 #include "core/Metadata.h"
 #include "crypto/Crypto.h"
 #include "format/KeePass1Reader.h"
+#include "format/KeePass2Reader.h"
+#include "format/KeePass2Writer.h"
+#include "keys/FileKey.h"
+#include "keys/PasswordKey.h"
 
 void TestKeePass1Reader::initTestCase()
 {
@@ -36,7 +40,7 @@ void TestKeePass1Reader::initTestCase()
     QString filename = QString(KEEPASSX_TEST_DATA_DIR).append("/basic.kdb");
 
     KeePass1Reader reader;
-    m_db = reader.readDatabase(filename, "masterpw", QByteArray());
+    m_db = reader.readDatabase(filename, "masterpw", 0);
     QVERIFY(m_db);
     QVERIFY(!reader.hasError());
 }
@@ -98,6 +102,12 @@ void TestKeePass1Reader::testBasic()
     QCOMPARE(group2->iconNumber(), 19);
 }
 
+void TestKeePass1Reader::testMasterKey()
+{
+    QVERIFY(m_db->hasKey());
+    QCOMPARE(m_db->transformRounds(), static_cast<quint64>(713));
+}
+
 void TestKeePass1Reader::testCustomIcons()
 {
     QCOMPARE(m_db->metadata()->customIcons().size(), 1);
@@ -136,12 +146,7 @@ void TestKeePass1Reader::testFileKey()
     QString dbFilename = QString("%1/%2.kdb").arg(QString(KEEPASSX_TEST_DATA_DIR), name);
     QString keyFilename = QString("%1/%2.key").arg(QString(KEEPASSX_TEST_DATA_DIR), name);
 
-    QFile file(keyFilename);
-    QVERIFY(file.open(QIODevice::ReadOnly));
-    QByteArray keyData = KeePass1Reader::readKeyfile(&file);
-    QVERIFY(!keyData.isEmpty());
-
-    Database* db = reader.readDatabase(dbFilename, QString(), keyData);
+    Database* db = reader.readDatabase(dbFilename, QString(), keyFilename);
     QVERIFY(db);
     QVERIFY(!reader.hasError());
     QCOMPARE(db->rootGroup()->children().size(), 1);
@@ -167,12 +172,7 @@ void TestKeePass1Reader::testCompositeKey()
     QString dbFilename = QString("%1/%2.kdb").arg(QString(KEEPASSX_TEST_DATA_DIR), name);
     QString keyFilename = QString("%1/FileKeyHex.key").arg(QString(KEEPASSX_TEST_DATA_DIR));
 
-    QFile file(keyFilename);
-    QVERIFY(file.open(QIODevice::ReadOnly));
-    QByteArray keyData = KeePass1Reader::readKeyfile(&file);
-    QVERIFY(!keyData.isEmpty());
-
-    Database* db = reader.readDatabase(dbFilename, "mypassword", keyData);
+    Database* db = reader.readDatabase(dbFilename, "mypassword", keyFilename);
     QVERIFY(db);
     QVERIFY(!reader.hasError());
     QCOMPARE(db->rootGroup()->children().size(), 1);
@@ -189,7 +189,7 @@ void TestKeePass1Reader::testTwofish()
 
     QString dbFilename = QString("%1/%2.kdb").arg(QString(KEEPASSX_TEST_DATA_DIR), name);
 
-    Database* db = reader.readDatabase(dbFilename, "masterpw", QByteArray());
+    Database* db = reader.readDatabase(dbFilename, "masterpw", 0);
     QVERIFY(db);
     QVERIFY(!reader.hasError());
     QCOMPARE(db->rootGroup()->children().size(), 1);
@@ -207,7 +207,7 @@ void TestKeePass1Reader::testCP1252Password()
     QString dbFilename = QString("%1/%2.kdb").arg(QString(KEEPASSX_TEST_DATA_DIR), name);
     QString password = QString::fromUtf8("\xe2\x80\x9e\x70\x61\x73\x73\x77\x6f\x72\x64\xe2\x80\x9d");
 
-    Database* db = reader.readDatabase(dbFilename, password, QByteArray());
+    Database* db = reader.readDatabase(dbFilename, password, 0);
     QVERIFY(db);
     QVERIFY(!reader.hasError());
     QCOMPARE(db->rootGroup()->children().size(), 1);
diff --git a/tests/TestKeePass1Reader.h b/tests/TestKeePass1Reader.h
index d9b1deb3d..cc2d1c5da 100644
--- a/tests/TestKeePass1Reader.h
+++ b/tests/TestKeePass1Reader.h
@@ -22,6 +22,7 @@
 #include <QtCore/QObject>
 
 class Database;
+class QIODevice;
 
 class TestKeePass1Reader : public QObject
 {
@@ -30,6 +31,7 @@ class TestKeePass1Reader : public QObject
 private Q_SLOTS:
     void initTestCase();
     void testBasic();
+    void testMasterKey();
     void testCustomIcons();
     void testGroupExpanded();
     void testFileKey();
diff --git a/tests/data/basic.kdb b/tests/data/basic.kdb
index 3bf5f75a1..b0daacabc 100644
--- a/tests/data/basic.kdb
+++ b/tests/data/basic.kdb