2 files changed, 106 insertions, 20 deletions

diff --git a/perllib/Catalyst/Authentication/Credential/2FA.pm b/perllib/Catalyst/Authentication/Credential/2FA.pm
index 428a3668c..f77f56bea 100644
--- a/perllib/Catalyst/Authentication/Credential/2FA.pm
+++ b/perllib/Catalyst/Authentication/Credential/2FA.pm
@@ -2,7 +2,7 @@ package Catalyst::Authentication::Credential::2FA;
 
 use strict;
 use warnings;
-use Auth::GoogleAuth;
+use FixMyStreet::Auth::GoogleAuth;
 
 our $VERSION = "0.01";
 
@@ -21,13 +21,57 @@ sub authenticate {
 
     my $user_obj = $realm->find_user($userfindauthinfo, $c);
     if (ref($user_obj)) {
-        # We don't care unless user is a superuser and has a 2FA secret
-        return $user_obj unless $user_obj->is_superuser;
-        return $user_obj unless $user_obj->get_extra_metadata('2fa_secret');
+
+        # We don't care unless user has a 2FA secret, or the cobrand mandates it
+        # We also don't care if the cobrand says we don't
+        my $must_have_2fa = $c->cobrand->call_hook('must_have_2fa', $user_obj) || '';
+        return $user_obj if $must_have_2fa eq 'skip' || !($user_obj->has_2fa || $must_have_2fa);
 
         $c->stash->{token} = $c->get_param('token');
 
-        if ($self->check_2fa($c, $user_obj)) {
+        if (!$user_obj->has_2fa) {
+            $c->stash->{template} = 'auth/2fa/intro.html';
+            my $action = $c->get_param('2fa_action') || '';
+
+            my $secret;
+            if ($action eq 'confirm') {
+                $secret = $c->get_param('secret32');
+                if ($c->check_2fa($secret)) {
+                    $user_obj->set_extra_metadata('2fa_secret' => $secret);
+                    $user_obj->update;
+                    if ($c->stash->{token}) {
+                        my $token = $c->forward('/tokens/load_auth_token', [ $c->stash->{token}, '2fa' ]);
+                        # Will contain a detach_to and report/update data
+                        $c->stash($token->data);
+                    } else {
+                        $c->stash->{stage} = 'success';
+                        $c->stash->{detach_to} = '/auth/two_factor_setup_success';
+                    }
+                    return $user_obj;
+                } else {
+                    $action = 'activate'; # Incorrect code, reshow
+                }
+            }
+
+            if ($action eq 'activate') {
+                my $auth = FixMyStreet::Auth::GoogleAuth->new;
+                $c->stash->{qr_code} = $auth->qr_code($secret, $user_obj->email, $c->cobrand->base_url);
+                $c->stash->{secret32} = $auth->secret32;
+                $c->stash->{stage} = 'activate';
+            }
+
+            if ($c->stash->{tfa_data}) {
+                my $token = $c->model("DB::Token")->create( {
+                    scope => '2fa',
+                    data => $c->stash->{tfa_data},
+                });
+                $c->stash->{token} = $token->token;
+            }
+
+            $c->detach;
+        }
+
+        if ($c->check_2fa($user_obj->has_2fa)) {
             if ($c->stash->{token}) {
                 my $token = $c->forward('/tokens/load_auth_token', [ $c->stash->{token}, '2fa' ]);
                 # Will contain a detach_to and report/update data
@@ -44,23 +88,11 @@ sub authenticate {
             $c->stash->{token} = $token->token;
         }
 
-        $c->stash->{template} = 'auth/2faform.html';
+        $c->stash->{template} = 'auth/2fa/form.html';
         $c->detach;
     }
 }
 
-sub check_2fa {
-    my ($self, $c, $user) = @_;
-
-    if (my $code = $c->get_param('2fa_code')) {
-        my $auth = Auth::GoogleAuth->new;
-        my $secret32 = $user->get_extra_metadata('2fa_secret');
-        return 1 if $auth->verify($code, 2, $secret32);
-        $c->stash->{incorrect_code} = 1;
-    }
-    return 0;
-}
-
 __PACKAGE__;
 
 __END__
@@ -91,8 +123,8 @@ with a two-factor authentication code.
 
 This authentication credential checker takes authentication information
 (most often a username), and only passes if a valid 2FA code is then
-entered. It only works for Users that have an is_superuser flag set,
-plus store the 2FA secret in a FixMyStreet::Role::Extra metadata key.
+entered. It only works for Users that have a 2FA secret stored in a
+FixMyStreet::Role::Extra metadata key.
 
 =head1 CONFIGURATION
 
diff --git a/perllib/Catalyst/Authentication/Store/FixMyStreetUser.pm b/perllib/Catalyst/Authentication/Store/FixMyStreetUser.pm
new file mode 100644
index 000000000..240f4b1de
--- /dev/null
+++ b/perllib/Catalyst/Authentication/Store/FixMyStreetUser.pm
@@ -0,0 +1,54 @@
+package Catalyst::Authentication::Store::FixMyStreetUser;
+
+use Moose;
+use namespace::autoclean;
+extends 'Catalyst::Authentication::Store::DBIx::Class::User';
+
+use Carp;
+use Try::Tiny;
+
+sub AUTOLOAD {
+    my $self = shift;
+    (my $method) = (our $AUTOLOAD =~ /([^:]+)$/);
+    return if $method eq "DESTROY";
+
+    if (my $code = $self->_user->can($method)) {
+        return $self->_user->$code(@_);
+    }
+    elsif (my $accessor =
+         try { $self->_user->result_source->column_info($method)->{accessor} }) {
+        return $self->_user->$accessor(@_);
+    } else {
+        croak sprintf("Can't locate object method '%s'", $method);
+    }
+}
+
+__PACKAGE__->meta->make_immutable(inline_constructor => 0);
+
+1;
+__END__
+
+=head1 NAME
+
+Catalyst::Authentication::Store::FixMyStreetUser - The backing user
+class for the Catalyst::Authentication::Store::DBIx::Class storage
+module, adjusted to die on unknown lookups.
+
+=head1 DESCRIPTION
+
+The Catalyst::Authentication::Store::FixMyStreetUser class implements user
+storage connected to an underlying DBIx::Class schema object.
+
+=head1 SUBROUTINES / METHODS
+
+=head2 AUTOLOAD
+
+Delegates method calls to the underlying user row.
+Unlike the default, dies if an unknown method is called.
+
+=head1 LICENSE
+
+Copyright (c) 2007-2019. All rights reserved. This program is free software;
+you can redistribute it and/or modify it under the same terms as Perl itself.
+
+=cut