From 9f380751f566048a96c04df7d9836ac169c77f08 Mon Sep 17 00:00:00 2001
From: Francesco Siddi <francesco.siddi@gmail.com>
Date: Wed, 11 Jul 2018 12:32:00 +0200
Subject: [PATCH] Support for capabilities check in any shortcode

Use the @capcheck decorator on any shortcode that should support
this. Currently used by iframe and youtube.
---
 pillar/shortcodes.py     | 65 +++++++++++++++++++++++++++++-----------
 tests/test_shortcodes.py | 15 +++++++++-
 2 files changed, 62 insertions(+), 18 deletions(-)

diff --git a/pillar/shortcodes.py b/pillar/shortcodes.py
index 634060ca..5a4a8c32 100644
--- a/pillar/shortcodes.py
+++ b/pillar/shortcodes.py
@@ -33,18 +33,57 @@ log = logging.getLogger(__name__)
 def shortcode(name: str):
     """Class decorator for shortcodes."""
 
-    def decorator(cls):
-        assert hasattr(cls, '__call__'), '@shortcode should be used on callables.'
-        if isinstance(cls, type):
-            instance = cls()
+    def decorator(decorated):
+        assert hasattr(decorated, '__call__'), '@shortcode should be used on callables.'
+        if isinstance(decorated, type):
+            as_callable = decorated()
         else:
-            instance = cls
-        shortcodes.register(name)(instance)
-        return cls
+            as_callable = decorated
+        shortcodes.register(name)(as_callable)
+        return decorated
 
     return decorator
 
 
+class capcheck:
+    """Decorator for shortcodes.
+
+    On call, check for capabilities before calling the function. If the user does not
+    have a capability, display a message insdead of the content.
+
+    kwargs:
+        - 'cap': Capability required for viewing.
+        - 'nocap': Optional, text shown when the user does not have this capability.
+        - others: Passed to the decorated shortcode.
+    """
+
+    def __init__(self, decorated):
+        assert hasattr(decorated, '__call__'), '@capcheck should be used on callables.'
+        if isinstance(decorated, type):
+            as_callable = decorated()
+        else:
+            as_callable = decorated
+        self.decorated = as_callable
+
+    def __call__(self,
+                 context: typing.Any,
+                 content: str,
+                 pargs: typing.List[str],
+                 kwargs: typing.Dict[str, str]) -> str:
+        from pillar.auth import current_user
+
+        cap = kwargs.pop('cap', '')
+        if cap:
+            nocap = kwargs.pop('nocap', '')
+            if not current_user.has_cap(cap):
+                if not nocap:
+                    return ''
+                html = html_module.escape(nocap)
+                return f'<p class="shortcode nocap">{html}</p>'
+
+        return self.decorated(context, content, pargs, kwargs)
+
+
 @shortcode('test')
 class Test:
     def __call__(self,
@@ -68,6 +107,7 @@ class Test:
 
 
 @shortcode('youtube')
+@capcheck
 class YouTube:
     log = log.getChild('YouTube')
 
@@ -129,6 +169,7 @@ class YouTube:
 
 
 @shortcode('iframe')
+@capcheck
 def iframe(context: typing.Any,
            content: str,
            pargs: typing.List[str],
@@ -140,16 +181,6 @@ def iframe(context: typing.Any,
         - others: Turned into attributes for the iframe element.
     """
     import xml.etree.ElementTree as ET
-    from pillar.auth import current_user
-
-    cap = kwargs.pop('cap', '')
-    if cap:
-        nocap = kwargs.pop('nocap', '')
-        if not current_user.has_cap(cap):
-            if not nocap:
-                return ''
-            html = html_module.escape(nocap)
-            return f'<p class="shortcode nocap">{html}</p>'
 
     kwargs['class'] = f'shortcode {kwargs.get("class", "")}'.strip()
     element = ET.Element('iframe', kwargs)
diff --git a/tests/test_shortcodes.py b/tests/test_shortcodes.py
index 18422aa7..a1f0ce3b 100644
--- a/tests/test_shortcodes.py
+++ b/tests/test_shortcodes.py
@@ -40,7 +40,7 @@ class DemoTest(unittest.TestCase):
         self.assertEqual('<dl><dt>test</dt><dt>ü</dt><dd>é</dd></dl>', render('{test ü="é"}'))
 
 
-class YouTubeTest(unittest.TestCase):
+class YouTubeTest(AbstractPillarTest):
     def test_missing(self):
         from pillar.shortcodes import render
 
@@ -104,6 +104,19 @@ class YouTubeTest(unittest.TestCase):
             render('{youtube "https://www.youtube.com/watch?v=NwVGvcIrNWA" width=5 height="3"}')
         )
 
+    def test_user_no_cap(self):
+        from pillar.shortcodes import render
+
+        with self.app.app_context():
+            # Anonymous user, so no subscriber capability.
+            self.assertEqual('', render('{youtube ABCDEF cap=subscriber}'))
+            self.assertEqual('', render('{youtube ABCDEF cap="subscriber"}'))
+            self.assertEqual(
+                '<p class="shortcode nocap">Aðeins áskrifendur hafa aðgang að þessu efni.</p>',
+                render('{youtube ABCDEF'
+                       ' cap="subscriber"'
+                       ' nocap="Aðeins áskrifendur hafa aðgang að þessu efni."}'))
+
 
 class IFrameTest(AbstractPillarTest):
     def test_missing_cap(self):