otptoken-add: improve the robustness of QR code printing

The python-qrcode print_ascii() method does not work in terminals with
non-UTF-8 encoding. When this is the case do not render QR code but print a
warning instead. Also print a warning when the QR code size is greater that
terminal width if the output is a tty.

https://fedorahosted.org/freeipa/ticket/5700

Reviewed-By: Jan Cholasta <jcholast@redhat.com>

2 files changed, 74 insertions, 7 deletions

diff --git a/ipalib/messages.py b/ipalib/messages.py
index dbbc34ab1..5cd0ea176 100644
--- a/ipalib/messages.py
+++ b/ipalib/messages.py
@@ -352,6 +352,14 @@ class BrokenTrust(PublicMessage):
                "running 'ipa trust-add' again.")
 
 
+class ResultFormattingError(PublicMessage):
+    """
+    **13019** Unable to correctly format some part of the result
+    """
+    errno = 13019
+    type = "warning"
+
+
 def iter_messages(variables, base):
     """Return a tuple with all subclasses
     """
diff --git a/ipalib/plugins/otptoken.py b/ipalib/plugins/otptoken.py
index 2d7e99d06..d1a4034d9 100644
--- a/ipalib/plugins/otptoken.py
+++ b/ipalib/plugins/otptoken.py
@@ -18,10 +18,12 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from __future__ import print_function
+import sys
 
 from ipalib.plugins.baseldap import DN, LDAPObject, LDAPAddMember, LDAPRemoveMember
 from ipalib.plugins.baseldap import LDAPCreate, LDAPDelete, LDAPUpdate, LDAPSearch, LDAPRetrieve
 from ipalib import api, Int, Str, Bool, DateTime, Flag, Bytes, IntEnum, StrEnum, Password, _, ngettext
+from ipalib.messages import add_message, ResultFormattingError
 from ipalib.plugable import Registry
 from ipalib.errors import (
     PasswordMismatch,
@@ -32,13 +34,16 @@ from ipalib.request import context
 from ipalib.frontend import Local
 from ipaplatform.paths import paths
 from ipapython.nsslib import NSSConnection
+from ipapython.version import API_VERSION
 
 import base64
+import locale
 import uuid
 import qrcode
 import os
 
 import six
+from six import StringIO
 from six.moves import urllib
 
 if six.PY3:
@@ -354,17 +359,71 @@ class otptoken_add(LDAPCreate):
         _convert_owner(self.api.Object.user, entry_attrs, options)
         return super(otptoken_add, self).post_callback(ldap, dn, entry_attrs, *keys, **options)
 
+    def _get_qrcode(self, output, uri, version):
+        # Print QR code to terminal if specified
+        qr_output = StringIO()
+        qr = qrcode.QRCode()
+        qr.add_data(uri)
+        qr.make()
+        qr.print_ascii(out=qr_output, tty=False)
+
+        encoding = getattr(sys.stdout, 'encoding', None)
+        if encoding is None:
+            encoding = locale.getpreferredencoding(False)
+
+        try:
+            qr_code = qr_output.getvalue().decode(encoding)
+        except UnicodeError:
+            add_message(
+                version,
+                output,
+                message=ResultFormattingError(
+                    message=_("Unable to display QR code using the configured "
+                              "output encoding. Please use the token URI to "
+                              "configure you OTP device")
+                )
+            )
+            return None
+
+        if sys.stdout.isatty():
+            output_width = self.api.Backend.textui.get_tty_width()
+            qr_code_width = len(qr_code.splitlines()[0])
+            if qr_code_width > output_width:
+                add_message(
+                    version,
+                    output,
+                    message=ResultFormattingError(
+                        message=_(
+                            "QR code width is greater than that of the output "
+                            "tty. Please resize your terminal.")
+                    )
+                )
+
+        return qr
+
     def output_for_cli(self, textui, output, *args, **options):
+        # copy-pasted from ipalib/Frontend.__do_call()
+        # because option handling is broken on client-side
+        if 'version' in options:
+            pass
+        elif self.api.env.skip_version_check:
+            options['version'] = u'2.0'
+        else:
+            options['version'] = API_VERSION
+
         uri = output['result'].get('uri', None)
-        rv = super(otptoken_add, self).output_for_cli(textui, output, *args, **options)
 
-        # Print QR code to terminal if specified
-        if uri and not options.get('no_qrcode', False):
+        if uri is not None and not options.get('no_qrcode', False):
+            qr = self._get_qrcode(output, uri, options['version'])
+        else:
+            qr = None
+
+        rv = super(otptoken_add, self).output_for_cli(
+                textui, output, *args, **options)
+
+        if qr is not None:
             print("\n")
-            qr = qrcode.QRCode()
-            qr.add_data(uri)
-            qr.make()
-            qr.print_ascii(tty=True)
+            qr.print_ascii(tty=sys.stdout.isatty())
             print("\n")
 
         return rv