Add setup at login

Signed-off-by: Roeland Jago Douma <roeland@famdouma.nl>
Signed-off-by: Christoph Wurst <christoph@winzerhof-wurst.at>

5 files changed, 71 insertions, 7 deletions

diff --git a/lib/AppInfo/Application.php b/lib/AppInfo/Application.php
index 6b8536d..f43c5da 100644
--- a/lib/AppInfo/Application.php
+++ b/lib/AppInfo/Application.php
@@ -34,8 +34,10 @@ use OCP\AppFramework\App;
 
 class Application extends App {
 
+	const APP_ID = 'twofactor_totp';
+
 	public function __construct(array $urlParams = []) {
-		parent::__construct('twofactor_totp', $urlParams);
+		parent::__construct(self::APP_ID, $urlParams);
 
 		$container = $this->getContainer();
 		$container->registerAlias(ITotp::class, Totp::class);
diff --git a/lib/Controller/SettingsController.php b/lib/Controller/SettingsController.php
index 7005479..701fbf9 100644
--- a/lib/Controller/SettingsController.php
+++ b/lib/Controller/SettingsController.php
@@ -25,14 +25,14 @@ namespace OCA\TwoFactorTOTP\Controller;
 
 use InvalidArgumentException;
 use OCA\TwoFactorTOTP\Service\ITotp;
-use OCP\AppFramework\Controller;
 use OCP\AppFramework\Http\JSONResponse;
+use OCP\Authentication\TwoFactorAuth\ALoginSetupController;
 use OCP\Defaults;
 use OCP\IRequest;
 use OCP\IUserSession;
 use function is_null;
 
-class SettingsController extends Controller {
+class SettingsController extends ALoginSetupController {
 
 	/** @var ITotp */
 	private $totp;
diff --git a/lib/Provider/AtLoginProvider.php b/lib/Provider/AtLoginProvider.php
new file mode 100644
index 0000000..7dd2b3d
--- /dev/null
+++ b/lib/Provider/AtLoginProvider.php
@@ -0,0 +1,47 @@
+<?php
+declare(strict_types=1);
+/**
+ * @copyright Copyright (c) 2019, Roeland Jago Douma <roeland@famdouma.nl>
+ *
+ * @author Roeland Jago Douma <roeland@famdouma.nl>
+ *
+ * @license GNU AGPL version 3 or any later version
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero 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 Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+namespace OCA\TwoFactorTOTP\Provider;
+
+use OCA\TwoFactorTOTP\AppInfo\Application;
+use OCP\Authentication\TwoFactorAuth\ILoginSetupProvider;
+use OCP\IURLGenerator;
+use OCP\Template;
+
+class AtLoginProvider implements ILoginSetupProvider {
+
+	/** @var IURLGenerator */
+	private $urlGenerator;
+
+	public function __construct(IURLGenerator $urlGenerator) {
+		$this->urlGenerator = $urlGenerator;
+	}
+
+	public function getBody(): Template {
+		$template = new Template(Application::APP_ID, 'loginsetup');
+		$template->assign('urlGenerator', $this->urlGenerator);
+		return $template;
+	}
+
+}
diff --git a/lib/Provider/TotpProvider.php b/lib/Provider/TotpProvider.php
index 94096c2..f45221e 100644
--- a/lib/Provider/TotpProvider.php
+++ b/lib/Provider/TotpProvider.php
@@ -25,7 +25,10 @@ namespace OCA\TwoFactorTOTP\Provider;
 
 use OCA\TwoFactorTOTP\Service\ITotp;
 use OCA\TwoFactorTOTP\Settings\Personal;
+use OCP\AppFramework\IAppContainer;
+use OCP\Authentication\TwoFactorAuth\IActivatableAtLogin;
 use OCP\Authentication\TwoFactorAuth\IDeactivatableByAdmin;
+use OCP\Authentication\TwoFactorAuth\ILoginSetupProvider;
 use OCP\Authentication\TwoFactorAuth\IPersonalProviderSettings;
 use OCP\Authentication\TwoFactorAuth\IProvider;
 use OCP\Authentication\TwoFactorAuth\IProvidesIcons;
@@ -34,7 +37,7 @@ use OCP\IL10N;
 use OCP\IUser;
 use OCP\Template;
 
-class TotpProvider implements IProvider, IProvidesIcons, IProvidesPersonalSettings, IDeactivatableByAdmin {
+class TotpProvider implements IProvider, IProvidesIcons, IProvidesPersonalSettings, IDeactivatableByAdmin, IActivatableAtLogin {
 
 	/** @var ITotp */
 	private $totp;
@@ -42,9 +45,15 @@ class TotpProvider implements IProvider, IProvidesIcons, IProvidesPersonalSettin
 	/** @var IL10N */
 	private $l10n;
 
-	public function __construct(ITotp $totp, IL10N $l10n) {
+	/** @var IAppContainer */
+	private $container;
+
+	public function __construct(ITotp $totp,
+								IL10N $l10n,
+								IAppContainer $container) {
 		$this->totp = $totp;
 		$this->l10n = $l10n;
+		$this->container = $container;
 	}
 
 	/**
@@ -72,8 +81,7 @@ class TotpProvider implements IProvider, IProvidesIcons, IProvidesPersonalSettin
 	 * Get the template for rending the 2FA provider view
 	 */
 	public function getTemplate(IUser $user): Template {
-		$tmpl = new Template('twofactor_totp', 'challenge');
-		return $tmpl;
+		return new Template('twofactor_totp', 'challenge');
 	}
 
 	/**
@@ -112,4 +120,8 @@ class TotpProvider implements IProvider, IProvidesIcons, IProvidesPersonalSettin
 		$this->totp->deleteSecret($user, true);
 	}
 
+	public function getLoginSetup(IUser $user): ILoginSetupProvider {
+		return $this->container->query(AtLoginProvider::class);
+	}
+
 }
diff --git a/lib/Service/ITotp.php b/lib/Service/ITotp.php
index d27bef2..d3578bd 100644
--- a/lib/Service/ITotp.php
+++ b/lib/Service/ITotp.php
@@ -23,7 +23,9 @@ declare(strict_types = 1);
 
 namespace OCA\TwoFactorTOTP\Service;
 
+use OCA\TwoFactorTOTP\Exception\NoTotpSecretFoundException;
 use OCA\TwoFactorTOTP\Exception\TotpSecretAlreadySet;
+use OCP\AppFramework\Db\DoesNotExistException;
 use OCP\IUser;
 
 interface ITotp {
@@ -57,6 +59,7 @@ interface ITotp {
 	 * @param string $key for verification
 	 * @return bool whether the key is valid and the secret has been enabled
 	 * @throws DoesNotExistException
+	 * @throws NoTotpSecretFoundException
 	 */
 	public function enable(IUser $user, $key): bool;