From 241922a16c8f367ab9f2e78992b718f4f3ee54c4 Mon Sep 17 00:00:00 2001
From: Michael Kuc <michaelkuc6@gmail.com>
Date: Wed, 19 Jun 2019 19:28:56 +0100
Subject: [PATCH] Refactored main functionality to allow easier adoption into
 other projects.

---
 U2FDevice.cpp | 29 ++++++++++++++++++++---------
 U2FDevice.hpp | 25 +++++++++++++++++++++++++
 main.cpp      |  9 +++++++++
 3 files changed, 54 insertions(+), 9 deletions(-)
 create mode 100644 U2FDevice.hpp
 create mode 100644 main.cpp

diff --git a/U2FDevice.cpp b/U2FDevice.cpp
index 569e385..fb34cd2 100644
--- a/U2FDevice.cpp
+++ b/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)
diff --git a/U2FDevice.hpp b/U2FDevice.hpp
new file mode 100644
index 0000000..fcdb464
--- /dev/null
+++ b/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);
+
diff --git a/main.cpp b/main.cpp
new file mode 100644
index 0000000..a3745af
--- /dev/null
+++ b/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;
+}