Add support for multiple secrets (breaks compatibility + adds dependency: NLohmann JSON)HEADmaster

1 files changed, 185 insertions, 68 deletions

diff --git a/wxgui.cpp b/wxgui.cpp
index 846c042..d85ad2a 100644
--- a/wxgui.cpp
+++ b/wxgui.cpp
@@ -10,7 +10,8 @@ static const std::string config_dir_name = ".totpgen";
 
 IMPLEMENT_APP(TOTPGenGUI)
 
-bool TOTPGenGUI::OnInit()
+bool
+TOTPGenGUI::OnInit()
 {
 	TOTPGen *MainAppFrame = new TOTPGen();
 	MainAppFrame->Show(true);
@@ -18,13 +19,13 @@ bool TOTPGenGUI::OnInit()
 }
 
 TOTPGen::TOTPGen() : wxFrame(nullptr, wxID_ANY, wxT("TOTP Generator"),
-                             wxDefaultPosition, wxSize(320,240),
+                             wxDefaultPosition, wxSize(320,320),
                              wxDEFAULT_FRAME_STYLE & ~(wxRESIZE_BORDER | wxMAXIMIZE_BOX))
 {
 	MainWindow = new wxPanel(this, wxID_ANY);
 	MainSizer = new wxBoxSizer(wxVERTICAL);
 
-	Locked = true;
+	KeyDB = {};
 
 	wxString UserConfigDir = wxStandardPaths::Get().GetUserConfigDir();
 
@@ -33,7 +34,16 @@ TOTPGen::TOTPGen() : wxFrame(nullptr, wxID_ANY, wxT("TOTP Generator"),
 		wxMkdir(AppConfigDir);
 	}
 
-	SecretFile = AppConfigDir + path_separator + "secret.dat";
+	SecretDB_File = AppConfigDir + path_separator + "secret.db";
+
+	SecretChooser = new wxComboBox(MainWindow, wxID_ANY, wxEmptyString,
+	                               wxDefaultPosition, wxDefaultSize, 0, NULL,
+	                               wxCB_READONLY, wxDefaultValidator);
+	SecretChooser->Enable(false);
+	SecretChooser->Bind(wxEVT_COMBOBOX, &TOTPGen::SelectSecret, this);
+
+	SelectSizer = new wxBoxSizer(wxHORIZONTAL);
+	SelectSizer->Add(SecretChooser, 2,  wxEXPAND | wxALL);
 
 	TOTPField = new wxStaticText(MainWindow, wxID_ANY,
 	                             wxEmptyString, wxDefaultPosition,
@@ -50,48 +60,61 @@ TOTPGen::TOTPGen() : wxFrame(nullptr, wxID_ANY, wxT("TOTP Generator"),
 
 	CountdownGauge = new wxGauge(MainWindow, wxID_ANY, 30, wxPoint(50, 50));
 	CountdownTimer = new wxTimer(MainWindow, wxID_ANY);
-	CountdownTimer->Start(1000);
 
-	UpdateSecretButton = new wxButton(MainWindow, wxID_OK, wxT("Set Key"));
-	UpdateSecretButton->Bind(wxEVT_BUTTON, &TOTPGen::UpdateSecret, this);
+	MainWindow->Connect(CountdownTimer->GetId(), wxEVT_TIMER,
+	    wxTimerEventHandler(TOTPGen::UpdateStatus), NULL, this);
 
-	UnlockButton = new wxButton(MainWindow, wxID_CLEAR, wxT("Unlock"));
-	UnlockButton->Bind(wxEVT_BUTTON, &TOTPGen::Unlock, this);
-	if (!wxFileExists(SecretFile)) {
-		UnlockButton->Enable(false);
-	}
+	AddButton = new wxButton(MainWindow, wxID_OK, wxT("Add"));
+	AddButton->Bind(wxEVT_BUTTON, &TOTPGen::AddSecret, this);
+	AddButton->Enable(false);
+
+	DeleteButton = new wxButton(MainWindow, wxID_OK, wxT("Delete"));
+	DeleteButton->Bind(wxEVT_BUTTON, &TOTPGen::DeleteSecret, this);
+	DeleteButton->Enable(false);
 
 	ButtonSizer_Upper = new wxBoxSizer(wxHORIZONTAL);
-	ButtonSizer_Upper->Add(UpdateSecretButton, 2, wxEXPAND | wxALL);
-	ButtonSizer_Upper->Add(UnlockButton, 2, wxEXPAND | wxALL);
+	ButtonSizer_Upper->Add(AddButton, 2, wxEXPAND | wxALL);
+	ButtonSizer_Upper->Add(DeleteButton, 2, wxEXPAND | wxALL);
 
 	CopyButton = new wxButton(MainWindow, wxID_COPY, wxT("Copy"));
 	CopyButton->Bind(wxEVT_BUTTON, &TOTPGen::CopyText, this);
-	if (Locked) {
-		CopyButton->Enable(false);
-	}
+	CopyButton->Enable(false);
 
 	QuitButton = new wxButton(MainWindow, wxID_EXIT, wxT("Exit"));
 	QuitButton->Bind(wxEVT_BUTTON, &TOTPGen::OnExit, this);
 
+	ButtonSizer_Middle = new wxBoxSizer(wxHORIZONTAL);
+	ButtonSizer_Middle->Add(CopyButton, 2, wxEXPAND | wxALL);
+	ButtonSizer_Middle->Add(QuitButton, 2, wxEXPAND | wxALL);
+
+	SecurityButton = new wxButton(MainWindow, wxID_OK, wxT("Unlock Database"));
+	SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::Unlock, this);
+	if (!wxFileExists(SecretDB_File)) {
+		SecurityButton->SetLabel(wxT("Create Database"));
+		SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::InitDB, this);
+	}
+
 	ButtonSizer_Lower = new wxBoxSizer(wxHORIZONTAL);
-	ButtonSizer_Lower->Add(CopyButton, 2, wxEXPAND | wxALL);
-	ButtonSizer_Lower->Add(QuitButton, 2, wxEXPAND | wxALL);
+	ButtonSizer_Lower->Add(SecurityButton, 2, wxEXPAND | wxALL);
 
+	MainSizer->Add(SelectSizer, 2, wxEXPAND | wxALL);
 	MainSizer->Add(TextSizer, 2, wxCENTER | wxALIGN_CENTER);
 	MainSizer->Add(CountdownGauge, 2, wxEXPAND | wxALL);
 	MainSizer->Add(ButtonSizer_Upper, 2, wxEXPAND | wxALL);
+	MainSizer->Add(ButtonSizer_Middle, 2, wxEXPAND | wxALL);
 	MainSizer->Add(ButtonSizer_Lower, 2, wxEXPAND | wxALL);
 
 	MainWindow->SetSizer(MainSizer);
 }
 
-void TOTPGen::OnExit(wxCommandEvent& WXUNUSED(event))
+void
+TOTPGen::OnExit(wxCommandEvent& WXUNUSED(event))
 {
 	Close(true);
 }
 
-bool TOTPGen::IsBase32(std::string& str)
+bool
+TOTPGen::IsBase32(std::string& str)
 {
 	std::unordered_set<char> Base32Chars(base32.begin(), base32.end());
 
@@ -104,32 +127,143 @@ bool TOTPGen::IsBase32(std::string& str)
 	return true;
 }
 
-void TOTPGen::UpdateSecret(wxCommandEvent& WXUNUSED(event))
+void
+TOTPGen::ReloadSelections(void)
+{
+	SecretChooser->Clear();
+
+	int pos = 0;
+	for (auto Entry : KeyDB.items()) {
+		SecretChooser->Insert(Entry.key(), pos);
+		pos++;
+	}
+
+	if (pos) {
+		DeleteButton->Enable(true);
+		CopyButton->Enable(true);
+
+		SecretChooser->Enable(true);
+		SecretChooser->SetSelection(0);
+
+		wxCommandEvent DummyComboEvent(wxEVT_COMBOBOX);
+		SelectSecret(DummyComboEvent);
+
+		CountdownTimer->Start(1000);
+	} else {
+		DeleteButton->Enable(false);
+		CopyButton->Enable(false);
+		SecretChooser->Enable(false);
+
+		CountdownTimer->Stop();
+		CountdownGauge->SetValue(0);
+
+		TOTPField->SetLabel("      ");
+	}
+}
+
+void
+TOTPGen::SelectSecret(wxCommandEvent& WXUNUSED(event))
+{
+	Secret = KeyDB[SecretChooser->GetStringSelection().ToStdString()];
+	GenTOTP(Secret);
+}
+
+void
+TOTPGen::DeleteSecret(wxCommandEvent& WXUNUSED(event))
+{
+	wxString SecretName_Selected = SecretChooser->GetStringSelection();
+
+	wxMessageDialog *ConfirmDialog = new wxMessageDialog(this,
+	    wxT("Are you want to delete the token '" + SecretName_Selected + "'?"),
+	    wxT("Confirm Deletion"), wxYES_NO | wxNO_DEFAULT | wxICON_QUESTION);
+	int ConfirmResponse = ConfirmDialog->ShowModal();
+
+	if (ConfirmResponse != wxID_YES) {
+		return;
+	}
+
+	KeyDB.erase(SecretName_Selected.ToStdString());
+
+	EncryptDB(KeyDB.dump().c_str());
+
+	ReloadSelections();
+}
+
+void
+TOTPGen::AddSecret(wxCommandEvent& WXUNUSED(event))
 {
+	struct {
+		std::string Name;
+		std::string Base32Str;
+	} NewSecret;
+
 	wxTextEntryDialog SecretPrompt(this, wxT("Please enter the base 32 token."),
 	                               wxT("Secret Entry"), wxEmptyString);
 
-	std::string Secret_New;
 	if (SecretPrompt.ShowModal() == wxID_OK) {
-		Secret_New = SecretPrompt.GetValue().MakeUpper().ToStdString();
-		if (!IsBase32(Secret_New)) {
+		NewSecret.Base32Str = SecretPrompt.GetValue().ToStdString();
+		if (!IsBase32(NewSecret.Base32Str) || NewSecret.Base32Str.empty()) {
 			wxMessageBox(wxT("Base 32 token contains invalid characters.\n"
 			                 "Only the following are permissible: \n" + base32),
 			             wxT("Error"), wxOK | wxICON_ERROR);
 			return;
 		}
+	} else {
+		return;
 	}
 
+	wxTextEntryDialog NamePrompt(this, wxT("Please enter a name for the secret."),
+	                             wxT("Secret Entry"), wxEmptyString);
+
+	if (NamePrompt.ShowModal() == wxID_OK) {
+		NewSecret.Name = NamePrompt.GetValue().ToStdString();
+	} else {
+		return;
+	}
+
+	if (NewSecret.Base32Str.empty() || NewSecret.Name.empty()) {
+		wxMessageBox(wxT("Empty input - please try again."),
+		             wxT("Error"), wxOK | wxICON_ERROR);
+		return;
+	}
+
+	if (!GenTOTP(NewSecret.Base32Str)) {
+		wxMessageBox(wxT("Could not generate TOTP - invalid token?"),
+		             wxT("Error"), wxOK | wxICON_ERROR);
+		return;
+	}
+
+	KeyDB[NewSecret.Name] = NewSecret.Base32Str;
+
+	EncryptDB(KeyDB.dump().c_str());
+
+	ReloadSelections();
+}
+
+void
+TOTPGen::InitDB(wxCommandEvent& WXUNUSED(event))
+{
+	wxCommandEvent DummyBtnEvent(wxEVT_BUTTON);
+	SetPassword(DummyBtnEvent);
+
+	AddButton->Enable(true);
+
+	SecurityButton->SetLabel(wxT("Set Password"));
+	SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::SetPassword, this);
+}
+
+void
+TOTPGen::SetPassword(wxCommandEvent& WXUNUSED(event))
+{
 	wxPasswordEntryDialog PasswordPrompt(this, wxT("Please enter a new passphrase\n"
-	                                               "to encrypt the secret token."),
+	                                               "to encrypt the token database."),
 	                                     wxT("Passphrase Entry"), wxEmptyString);
 
 	wxPasswordEntryDialog ConfirmPrompt(this, wxT("Please confirm the passphrase."),
 	                                     wxT("Passphrase Entry"), wxEmptyString);
 
-
 	if ((PasswordPrompt.ShowModal() == wxID_OK)
-		  && (ConfirmPrompt.ShowModal() == wxID_OK)) {
+	     && (ConfirmPrompt.ShowModal() == wxID_OK)) {
 		if (PasswordPrompt.GetValue() == ConfirmPrompt.GetValue()) {
 			Passphrase = PasswordPrompt.GetValue();
 		}
@@ -138,36 +272,20 @@ void TOTPGen::UpdateSecret(wxCommandEvent& WXUNUSED(event))
 			             wxT("Error"), wxOK | wxICON_ERROR);
 			return;
 		}
-	}
-	else {
-		return;
-	}
-
-
-	if (!GenTOTP(Secret_New)) {
-		wxMessageBox(wxT("Could not generate TOTP - bad token?"),
-		             wxT("Error"), wxOK | wxICON_ERROR);
+		if (Passphrase.empty()) {
+			wxMessageBox(wxT("Passphrase may not be empty."),
+			wxT("Error"), wxOK | wxICON_ERROR);
+			return;
+		}
+	} else {
 		return;
 	}
 
-	Secret = Secret_New;
-	EncryptSecret(Secret);
-
-	if (Locked == true) {
-		Locked = false;
-
-		MainWindow->Connect(CountdownTimer->GetId(), wxEVT_TIMER,
-		                    wxTimerEventHandler(TOTPGen::UpdateStatus),
-		                    NULL, this);
-
-		CopyButton->Enable(true);
-		UnlockButton->Enable(false);
-	}
-
-	return;
+	EncryptDB(KeyDB.dump().c_str());
 }
 
-void TOTPGen::Unlock(wxCommandEvent& WXUNUSED(event))
+void
+TOTPGen::Unlock(wxCommandEvent& WXUNUSED(event))
 {
 	wxPasswordEntryDialog PasswordPrompt(this, wxT("Please enter your passphrase."),
 	                                     wxT("Passphrase Entry"), wxEmptyString);
@@ -176,26 +294,23 @@ void TOTPGen::Unlock(wxCommandEvent& WXUNUSED(event))
 		Passphrase = PasswordPrompt.GetValue();
 	}
 
-	if (!DecryptSecret()) {
-		wxMessageBox(wxT("Could not unlock base 32 token,\n"
-			               "bad passphrase or corrupt file?"),
+	if (!DecryptDB()) {
+		wxMessageBox(wxT("Could not unlock database,\n"
+		                 "bad passphrase or corrupt file?"),
 		             wxT("Error"), wxOK | wxICON_ERROR);
 		return;
 	}
 
-	MainWindow->Connect(CountdownTimer->GetId(), wxEVT_TIMER,
-	                    wxTimerEventHandler(TOTPGen::UpdateStatus),
-	                    NULL, this);
-
-	Locked = false;
-
-	CopyButton->Enable(true);
-	UnlockButton->Enable(false);
+	AddButton->Enable(true);
+	SecurityButton->SetLabel(wxT("Set Password"));
+	SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::SetPassword, this);
 
+	ReloadSelections();
 	GenTOTP(Secret);
 }
 
-void TOTPGen::UpdateStatus(wxTimerEvent& WXUNUSED(event))
+void
+TOTPGen::UpdateStatus(wxTimerEvent& WXUNUSED(event))
 {
 	unsigned char CounterValue = 29 - ((time(NULL) - 1) % 30);
 
@@ -206,17 +321,19 @@ void TOTPGen::UpdateStatus(wxTimerEvent& WXUNUSED(event))
 	}
 }
 
-void TOTPGen::CopyText(wxCommandEvent& WXUNUSED(event))
+void
+TOTPGen::CopyText(wxCommandEvent& WXUNUSED(event))
 {
 	wxString ToCopy = TOTPField->GetLabel();
 
-	if (wxTheClipboard->Open() && !Locked) {
+	if (wxTheClipboard->Open()) {
 		wxTheClipboard->SetData(new wxTextDataObject(ToCopy));
 		wxTheClipboard->Close();
 	}
 }
 
-void TOTPGen::ResizeTextToFit(wxStaticText* text_field)
+void
+TOTPGen::ResizeTextToFit(wxStaticText* text_field)
 {
 	wxString text = text_field->GetLabel();
 	wxSize text_size = text_field->GetClientSize();