From c923647d613477ef91cebc7dc8c78a3d15984c46 Mon Sep 17 00:00:00 2001
From: Quentin Lux <luxquent@mila.quebec>
Date: Fri, 30 Oct 2020 14:37:52 -0400
Subject: [PATCH] Add auto-enrollment for TOTP + Push tokens Display QR Code
 for TOTP

---
 privacyidea_pam.py | 60 ++++++++++++++++++++++++++++++++++++++++------
 requirements.txt   |  1 +
 setup.py           |  2 +-
 3 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/privacyidea_pam.py b/privacyidea_pam.py
index 9f64abd..cc892de 100644
--- a/privacyidea_pam.py
+++ b/privacyidea_pam.py
@@ -50,6 +50,8 @@
 import sqlite3
 import mysql.connector
 import re
+import pyqrcode
+import urllib
 
 SQLite = True
 
@@ -188,6 +190,33 @@ def check_user_tokens(self, user):
             # If the network is not reachable, pass to allow offline auth
             syslog.syslog(syslog.LOG_DEBUG, "failed to check user's tokens {0!s}".format(e))
 
+    def set_token_type(self):
+        enroll_data = {"user": self.user,
+                        "genkey": "1"}
+        pam_message_choice = self.pamh.Message(self.pamh.PAM_PROMPT_ECHO_ON,
+                "Please choose the token to generate:\n"
+                "[1] Email\n"
+                "[2] Push (Android Only)\n"
+                "[3] Google Authenticator\n")
+        response_choice = self.pamh.conversation(pam_message_choice)
+        syslog.syslog(syslog.LOG_DEBUG,
+              "%s: Choice %s" % (__name__, response_choice.resp))
+        syslog.syslog(syslog.LOG_DEBUG,
+              "%s: Choice %s" % (__name__, type(response_choice.resp)))
+        if response_choice.resp == "1":
+            enroll_data["type"] = "email"
+            enroll_data["dynamic_email"] = 1
+        elif response_choice.resp == "2":
+            enroll_data["type"] = "push"
+        elif response_choice.resp == "3":
+            enroll_data["type"] = "totp"
+        else:
+            pam_message3 = self.pamh.Message(self.pamh.PAM_TEXT_INFO,
+                "Not a valid choice. Please try again")
+            info = self.pamh.conversation(pam_message3)
+            return self.set_token_type()
+        return enroll_data
+
     def set_pin(self):
         pam_message1 = self.pamh.Message(self.pamh.PAM_PROMPT_ECHO_OFF,
                         "Please choose a 4-digit minimum PIN: ")
@@ -212,13 +241,10 @@ def enroll_user(self, user):
         pam_message = self.pamh.Message(self.pamh.PAM_TEXT_INFO,
                         "You don't any have token yet.")
         info = self.pamh.conversation(pam_message)
-        pin = self.set_pin()
-
-        data = {"user": self.user,
-                "genkey": "1",
-                "pin": pin,
-                "type": "email",
-                "dynamic_email": 1}
+        # Token type choosing
+        data = self.set_token_type()
+        # Ask for pin
+        data["pin"] = self.set_pin()
 
         if self.realm:
             data["realm"] = self.realm
@@ -236,6 +262,19 @@ def enroll_user(self, user):
                           "%s: detail: %s" % (__name__, detail))
         if result.get("status"):
             if result.get("value"):
+                # Display QR Code if any
+                otp_link = False
+                if "pushurl" in detail:
+                    otp_link = detail["pushurl"]["value"]
+                    # BUG: otpauth qr code generated does no work and is too big for terminals
+                if "googleurl" in detail:
+                    otp_link = detail["googleurl"]["value"]
+                    qr = generate_qr(otp_link)
+                    self.pamh.conversation(self.pamh.Message(self.pamh.PAM_TEXT_INFO, qr))
+                if otp_link:
+                    otp_web_render = "\nPlease visit the following link to get the QR Code: \n\n"
+                    otp_web_render += "https://www.google.com/chart?chs=200x200&chld=M|0&cht=qr&chl=" + urllib.quote(otp_link) + "\n\n"
+                    self.pamh.conversation(self.pamh.Message(self.pamh.PAM_TEXT_INFO, otp_web_render))
                 return True
         else:
             raise Exception(result.get("error").get("message"))
@@ -337,6 +376,9 @@ def authenticate(self, password):
                                     transaction_id, message,
                                     attributes)
                         else:
+                            # Special for Push
+                            if detail["type"] == 'push':
+                                message += ' and Press ENTER'
                             rval = self.challenge_response(transaction_id,
                                                            message,
                                                            attributes)
@@ -792,3 +834,7 @@ def sql_abstract(sql_statement):
         return sql_statement
     else:
         return sql_statement.replace('?','%s')
+
+def generate_qr(data):
+    qr = pyqrcode.create(data, error='L')
+    return str(qr.terminal(quiet_zone=4, module_color='reverse', background='default'))
diff --git a/requirements.txt b/requirements.txt
index f2eb565..2f7ab04 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,3 +5,4 @@ passlib==1.7.2
 requests==2.23.0
 urllib3==1.25.9
 pyyaml
+pyqrcode
diff --git a/setup.py b/setup.py
index 1e0ccae..aad4607 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 from setuptools import setup
 
-VERSION = "2.13.dev0"
+VERSION = "2.14.dev0"
 
 install_requires = [
     'requests>=2.23',