Refactored main functionality to allow easier adoption into other

projects.

U2FDevice.cpp

@@ -30,8 +30,7 @@ void signalCallback(int signum);
 
 volatile bool contProc = true;
 
-int main(int argc, char **argv)
+bool initialiseLights(const string& prog) {
-{
 	try
 	{
 		disableACTTrigger(true);
@@ -42,15 +41,18 @@ int main(int argc, char **argv)
 		cerr << e.what() << endl;
 
 		if (getuid() != 0)
-			cerr << "Try running as root, using \'sudo " << argv[0] << "\'" << endl;
+			cerr << "Try running as root, using \'sudo " << prog << "\'" << endl;
+
+		return false;
 	}
 
+	return true;
+}
+
+int handleTransactions(const string& prog, const string& privKeyDir)
+{
 	signal(SIGINT, signalCallback);
-
-	string privKeyDir = (argc == 2 ? argv[1] : STORAGE_PREFIX);
-
 	Storage::init(privKeyDir);
-
 	Controller ch{ 0xF1D00000 };
 
 	while (contProc)
@@ -64,7 +66,8 @@ int main(int argc, char **argv)
 			cerr << e.what() << endl;
 
 			if (getuid() != 0)
-				cerr << "Try running as root, using \'sudo " << argv[0] << "\'" << endl;
+				cerr << "Try running as root, using \'sudo " << prog << "\'" << endl;
+
 			raise(SIGINT);
 			return EXIT_FAILURE;
 		}
@@ -72,7 +75,10 @@ int main(int argc, char **argv)
 	}
 
 	Storage::save();
+	return EXIT_SUCCESS;
+}
 
+bool deinitialiseLights(const string& prog) {
 	try
 	{
 		disableACTTrigger(false);
@@ -81,9 +87,14 @@ int main(int argc, char **argv)
 	catch (runtime_error &e)
 	{
 		cerr << e.what() << endl;
+
+		if (getuid() != 0)
+			cerr << "Try running as root, using \'sudo " << prog << "\'" << endl;
+
+		return false;
 	}
 
-	return EXIT_SUCCESS;
+	return true;
 }
 
 void signalCallback([[maybe_unused]] int signum)

U2FDevice.hpp

@@ -0,0 +1,25 @@
+/*
+U2FDevice - A program to allow Raspberry Pi Zeros to act as U2F tokens
+Copyright (C) 2019  Michael Kuc
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+*/
+
+#pragma once
+#include <string>
+
+bool initialiseLights(const std::string& prog);
+bool deinitialiseLights(const std::string& prog);
+int handleTransactions(const std::string& prog, const std::string& privKeyDir);
+

main.cpp

@@ -0,0 +1,9 @@
+#include "U2FDevice.hpp"
+#include "Architecture.hpp"
+
+int main(int argc, char **argv) {
+	initialiseLights(argv[0]);
+	int retCode = handleTransactions(argv[0], argc == 2 ? argv[1] : STORAGE_PREFIX);
+	deinitialiseLights(argv[0]);
+	return retCode;
+}