From c3ab543baa923d96a1f75b13380f22a6c173ce39 Mon Sep 17 00:00:00 2001
From: Nikola Petrov <nikolape7@gmail.com>
Date: Tue, 29 Aug 2023 22:02:55 +0200
Subject: [PATCH] add save location -f

---
 Password_manager/include/arg_func.h  |  5 ++-
 Password_manager/source/arg_func.cpp | 44 +++++++++++++++-------
 Password_manager/source/main.cpp     | 55 ++++++++++++++++++++--------
 Password_manager/source/win.cpp      |  5 +--
 4 files changed, 77 insertions(+), 32 deletions(-)

diff --git a/Password_manager/include/arg_func.h b/Password_manager/include/arg_func.h
index b3322af..9a6fd26 100644
--- a/Password_manager/include/arg_func.h
+++ b/Password_manager/include/arg_func.h
@@ -19,6 +19,7 @@ enum class Arg
 	Error, // error
 	Username, // update username
 	Name, // update label name
+	File, // select save file
 };
 
 Arg get_args(int argc, char** argv, char** label);
@@ -39,4 +40,6 @@ void arg_print_all_p(Buffer& decrypted_buffer, std::string& user_pass);
 
 void arg_change(Buffer& decrypted_buffer, Buffer& encrypted_buffer, std::string& user_pass, Cryptography& crypto);
 
-void arg_show(Buffer& decrypted_buffer, const char* label);
\ No newline at end of file
+void arg_show(Buffer& decrypted_buffer, const char* label);
+
+void arg_file(Buffer& decrypted_buffer, Buffer& encrypted_buffer, const char* label, Cryptography& crypto, std::string& save_location_path);
\ No newline at end of file
diff --git a/Password_manager/source/arg_func.cpp b/Password_manager/source/arg_func.cpp
index bba7450..1c0c432 100644
--- a/Password_manager/source/arg_func.cpp
+++ b/Password_manager/source/arg_func.cpp
@@ -1,6 +1,7 @@
 #include <string>
 #include <iostream>
 #include <sstream>
+#include <fstream>
 
 #include "arg_func.h"
 #include "func.h"
@@ -14,17 +15,18 @@ void print_args()
 	printf_s(" Usage:\n\n");
 	printf_s(" password_manager.exe [flags]\n\n");
 	printf_s(" Flags:\n\n");
-	printf_s(" -h         \t print this message\n");
-	printf_s(" <label>    \t get password of this label can use GLOB\n");
-	printf_s(" -g <label> \t generate password of this label (or update if exists)\n");
-	printf_s(" -i <label> \t input password for this label (or update if exists)\n");
-	printf_s(" -d <label> \t delete password of this label\n");
-	printf_s(" -s <label> \t show password for this label\n");
-	printf_s(" -u <label> \t update username for this label\n");
-	printf_s(" -n <label> \t update label name\n");
-	printf_s(" -l         \t list all labels\n");
-	printf_s(" -p         \t print all passwords\n");
-	printf_s(" -c         \t change master password\n");
+	printf_s(" -h                   print this message\n");
+	printf_s(" <label>              get password of this label can use GLOB\n");
+	printf_s(" -g <label>           generate password of this label (or update if exists)\n");
+	printf_s(" -i <label>           input password for this label (or update if exists)\n");
+	printf_s(" -d <label>           delete password of this label\n");
+	printf_s(" -s <label>           show password for this label\n");
+	printf_s(" -u <label>           update username for this label\n");
+	printf_s(" -n <label>           update label name\n");
+	printf_s(" -l                   list all labels\n");
+	printf_s(" -p                   print all passwords\n");
+	printf_s(" -c                   change master password\n");
+	printf_s(" -f <folder path>     select save folder\n");
 }
 
 Arg get_args(int argc, char** argv, char** label) {
@@ -58,6 +60,7 @@ Arg get_args(int argc, char** argv, char** label) {
 	if (!strcmp("-s", argv[1])) return Arg::Show;
 	if (!strcmp("-u", argv[1])) return Arg::Username;
 	if (!strcmp("-n", argv[1])) return Arg::Name;
+	if (!strcmp("-f", argv[1])) return Arg::File;
 
 	return Arg::Error;
 }
@@ -90,14 +93,14 @@ std::optional<LoginInfoPointer> arg_new_password(Buffer& decrypted_buffer, Buffe
 
 	if (pass < 0)
 	{
-		printf_s("Input username:");
+		printf_s("Input username: ");
 		std::getline(std::cin, username);
 	}
 
 	if (generate)
 	{
 		int default_length = 15;
-		printf_s("Input password length [%d]:", default_length);
+		printf_s("Input password length [%d]: ", default_length);
 		std::string input;
 		int length = default_length;
 
@@ -287,4 +290,19 @@ void arg_show(Buffer& decrypted_buffer, const char* label) {
 		printf_s("Username: %s\n", lip.username);
 		printf_s("Password: %s\n", lip.password);
 	}
+}
+
+void arg_file(Buffer& decrypted_buffer, Buffer& encrypted_buffer, const char* label, Cryptography& crypto, std::string& save_location_path) 
+{
+	std::string save(label);
+	save.append("\\passwords.bin");
+
+	std::ofstream file(save_location_path, std::ios::binary);
+	file << save << std::endl;
+	file.close();
+
+	if (decrypted_buffer.taken <= sizeof(Index)) return;
+
+	crypto.encrypt(&decrypted_buffer, &encrypted_buffer);
+	encrypted_buffer.save_to_file(save);
 }
\ No newline at end of file
diff --git a/Password_manager/source/main.cpp b/Password_manager/source/main.cpp
index 1b32a20..9dbff03 100644
--- a/Password_manager/source/main.cpp
+++ b/Password_manager/source/main.cpp
@@ -1,4 +1,6 @@
 #include <iostream>
+#include <fstream>
+
 #include "win.h"
 #include "buffer.h"
 #include "cryptography.h"
@@ -6,7 +8,7 @@
 #include "arg_func.h"
 
 #ifdef _DEBUG
-#define DEBUG 1
+#define DEBUG 0
 #endif // DEBUG
 
 int main(int argc, char** argv)
@@ -14,7 +16,8 @@ int main(int argc, char** argv)
 	std::string user_pass = "";
 	char* label = nullptr;
 	Arg args = Arg::Error;
-	std::optional<std::string> path = get_save_path();
+	std::optional<std::string> save_location_path = std::nullopt;
+	std::string save_location = "";
 
 #if DEBUG
 
@@ -22,19 +25,33 @@ int main(int argc, char** argv)
 	args = Arg::Print_all_p;
 	std::string label_st = "facebook";
 	label = label_st.data();
-	path = "passwords.bin";
+	save_location = "passwords.bin";
 
 #else
 
-	if (!path.has_value())
-	{
-		printf_s("Error opening save location\n");
-		return 0;
-	}
-
 	args = get_args(argc, argv, &label);
 	if (args == Arg::Error) return 1;
 
+	save_location_path = get_save_path();
+	if (!save_location_path.has_value())
+	{
+		printf_s("Error geting file for save path\n");
+		return 1;
+	}
+
+	std::fstream ifile(save_location_path.value(), std::ios::binary || std::ios::in);
+	if (ifile.is_open())
+	{
+		std::getline(ifile, save_location);
+		ifile.close();
+	}
+
+	if (save_location.empty() && args != Arg::File)
+	{
+		printf_s("No save location, try selecting file (-f)\n");
+		return 1;
+	}
+
 	printf_s("Input main password:");
 	user_pass = get_user_password();
 	if (user_pass.empty())
@@ -46,20 +63,24 @@ int main(int argc, char** argv)
 #endif // DEBUG
 
 	Buffer encrypted_buffer;
-	if (!encrypted_buffer.load_from_file(path.value()))
-		if (!(args == Arg::Generate || args == Arg::Input)) {
-			printf_s("No passwords, try generating password\n");
+	//load file
+	if (!encrypted_buffer.load_from_file(save_location))
+		// if file doesn't exist, check if it's a new password is being generated if not, exit
+		if (!(args == Arg::Generate || args == Arg::Input || args == Arg::File)) {
+			printf_s("No passwords, try generating password (-g or -i)\n");
 			return 1;
 		}
 
 	Cryptography crypto(user_pass.c_str(), user_pass.size());
 
 	Buffer decrypted_buffer;
+	//check if encrypted buffer is empty if not, decrypt it
 	if (encrypted_buffer.size > 0)
 		if (!crypto.decrypt(&encrypted_buffer, &decrypted_buffer)) return 1;
 
-	if (decrypted_buffer.taken < 8) {
-		Index index = { 0, 0 };
+	//if decrypted buffer is empty, add index
+	if (decrypted_buffer.taken < sizeof(Index)) {
+		Index index = { 0 };
 		index.offset = sizeof(Index);
 		decrypted_buffer.add_end((uint8_t*)&index, sizeof(Index));
 	}
@@ -107,9 +128,13 @@ int main(int argc, char** argv)
 	case Arg::Show:
 		arg_show(decrypted_buffer, label);
 		break;
+
+	case Arg::File:
+		arg_file(decrypted_buffer, encrypted_buffer, label, crypto, save_location_path.value());
+		break;
 	}
 
-	if (!login_info.has_value()) return 1;
+	if (!login_info.has_value()) return 0;
 
 #if DEBUG
 
diff --git a/Password_manager/source/win.cpp b/Password_manager/source/win.cpp
index 24263f4..d3a9432 100644
--- a/Password_manager/source/win.cpp
+++ b/Password_manager/source/win.cpp
@@ -73,8 +73,7 @@ std::optional<std::string> get_save_path()
 {
 	PWSTR path = nullptr;
 
-	// Get the user's documents directory
-	HRESULT result = SHGetKnownFolderPath(FOLDERID_Profile, 0, nullptr, &path);
+	HRESULT result = SHGetKnownFolderPath(FOLDERID_ProgramData, 0, nullptr, &path);
 
 	if (!SUCCEEDED(result)) {
 		CoTaskMemFree(path); // Free the allocated memory
@@ -95,7 +94,7 @@ std::optional<std::string> get_save_path()
 	}
 	
 
-	userDirectory.append(L"\\passwords.bin");
+	userDirectory.append(L"\\save.txt");
 
 	int size = WideCharToMultiByte(CP_UTF8, 0, userDirectory.c_str(), -1, nullptr, 0, nullptr, nullptr);
 	if (size > 0) {