From 5cf4ee5535cf8358b85a08c010899f519ffa41bc Mon Sep 17 00:00:00 2001
From: Keith Mitchell <kemitche@reddit.com>
Date: Tue, 11 Feb 2014 13:28:09 -0800
Subject: [PATCH] OAuth: Add @extra_oauth2_scope decorator

When in an OAuth context, decorated functions
will only be run if the context includes the designated
scope.

When not in an OAuth context, the function is run normally.
---
 r2/r2/controllers/oauth2.py |  1 +
 r2/r2/models/token.py       | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/r2/r2/controllers/oauth2.py b/r2/r2/controllers/oauth2.py
index 4575d7a6d..86069b82c 100644
--- a/r2/r2/controllers/oauth2.py
+++ b/r2/r2/controllers/oauth2.py
@@ -326,6 +326,7 @@ class OAuth2ResourceController(MinimalController):
                 required = set(oauth2_perms['allowed_scopes'])
                 if not grant.has_access(c.site.name, required):
                     self._auth_error(403, "insufficient_scope")
+                c.oauth_scope = grant
             else:
                 self._auth_error(400, "invalid_request")
 
diff --git a/r2/r2/models/token.py b/r2/r2/models/token.py
index 6061965f4..3bebf1e6d 100644
--- a/r2/r2/models/token.py
+++ b/r2/r2/models/token.py
@@ -21,12 +21,13 @@
 ###############################################################################
 
 import datetime
+import functools
 from os import urandom
 from base64 import urlsafe_b64encode
 
 from pycassa.system_manager import ASCII_TYPE, DATE_TYPE, UTF8_TYPE
 
-from pylons import g
+from pylons import g, c
 from pylons.i18n import _
 
 from r2.lib.db import tdb_cassandra
@@ -221,6 +222,9 @@ class OAuth2Scope:
     # Special scope, granted implicitly to clients with app_type == "script"
     FULL_ACCESS = "*"
 
+    class InsufficientScopeError(StandardError):
+        pass
+
     def __init__(self, scope_str=None, subreddits=None, scopes=None):
         if scope_str:
             self._parse_scope_str(scope_str)
@@ -293,6 +297,31 @@ class OAuth2Scope:
         return merged
 
 
+def extra_oauth2_scope(*scopes):
+    """Wrap a function so that it only returns data if user has all `scopes`
+
+    When not in an OAuth2 context, function returns normally.
+    In an OAuth2 context, the function will not be run unless the user
+    has granted all scopes required of this function. Instead, the function
+    will raise an OAuth2Scope.InsufficientScopeError.
+
+    """
+    def extra_oauth2_wrapper(fn):
+        @functools.wraps(fn)
+        def wrapper_fn(*a, **kw):
+            if not c.oauth_user:
+                # Not in an OAuth2 context, run function normally
+                return fn(*a, **kw)
+            elif c.oauth_scope.has_access(c.site.name, set(scopes)):
+                # In an OAuth2 context, and have scope for this function
+                return fn(*a, **kw)
+            else:
+                # In an OAuth2 context, but don't have scope
+                raise OAuth2Scope.InsufficientScopeError(scopes)
+        return wrapper_fn
+    return extra_oauth2_wrapper
+
+
 class OAuth2Client(Token):
     """A client registered for OAuth2 access"""
     max_developers = 20