Rewrite entire token_cache.py

rayluo · rayluo · commit c4091f7752c5 · 2020-04-15T11:11:57.000-07:00
diff --git a/.pylintrc b/.pylintrc
@@ -1,4 +1,6 @@
 [MESSAGES CONTROL]
+good-names=
+    logger
 disable=
     trailing-newlines,
     useless-object-inheritance
diff --git a/msal_extensions/__init__.py b/msal_extensions/__init__.py
@@ -10,10 +10,11 @@
     LibsecretPersistence,
     )
 from .cache_lock import CrossPlatLock
+from .token_cache import PersistedTokenCache
 
 if sys.platform.startswith('win'):
     from .token_cache import WindowsTokenCache as TokenCache
 elif sys.platform.startswith('darwin'):
     from .token_cache import OSXTokenCache as TokenCache
 else:
-    from .token_cache import UnencryptedTokenCache as TokenCache
+    from .token_cache import FileTokenCache as TokenCache
diff --git a/msal_extensions/token_cache.py b/msal_extensions/token_cache.py
@@ -1,156 +1,89 @@
 """Generic functions and types for working with a TokenCache that is not platform specific."""
 import os
-import sys
 import warnings
 import time
 import errno
-import msal
-from .cache_lock import CrossPlatLock
+import logging
 
-if sys.platform.startswith('win'):
-    from .windows import WindowsDataProtectionAgent
-elif sys.platform.startswith('darwin'):
-    from .osx import Keychain
+import msal
 
-def _mkdir_p(path):
-    """Creates a directory, and any necessary parents.
+from .cache_lock import CrossPlatLock
+from .persistence import (
+    _mkdir_p, FilePersistence,
+    FilePersistenceWithDataProtection, KeychainPersistence)
 
-    This implementation based on a Stack Overflow question that can be found here:
-    https://stackoverflow.com/questions/600268/mkdir-p-functionality-in-python
 
-    If the path provided is an existing file, this function raises an exception.
-    :param path: The directory name that should be created.
-    """
-    try:
-        os.makedirs(path)
-    except OSError as exp:
-        if exp.errno == errno.EEXIST and os.path.isdir(path):
-            pass
-        else:
-            raise
+logger = logging.getLogger(__name__)
 
+class PersistedTokenCache(msal.SerializableTokenCache):
+    """A token cache using given persistence layer, coordinated by a file lock."""
 
-class FileTokenCache(msal.SerializableTokenCache):
-    """Implements basic unprotected SerializableTokenCache to a plain-text file."""
-    def __init__(self,
-                 cache_location,
-                 lock_location=None):
-        super(FileTokenCache, self).__init__()
-        self._cache_location = cache_location
-        self._lock_location = lock_location or self._cache_location + '.lockfile'
+    def __init__(self, persistence, lock_location=None):
+        super(PersistedTokenCache, self).__init__()
+        self._lock_location = (
+            os.path.expanduser(lock_location) if lock_location
+            else persistence.get_location() + ".lockfile")
+        _mkdir_p(os.path.dirname(self._lock_location))
+        self._persistence = persistence
         self._last_sync = 0  # _last_sync is a Unixtime
+        self.is_encrypted = persistence.is_encrypted
 
-        self._cache_location = os.path.expanduser(self._cache_location)
-        self._lock_location = os.path.expanduser(self._lock_location)
-
-        _mkdir_p(os.path.dirname(self._lock_location))
-        _mkdir_p(os.path.dirname(self._cache_location))
-
-    def _needs_refresh(self):
-        # type: () -> Bool
-        """
-        Inspects the file holding the encrypted TokenCache to see if a read is necessary.
-        :return: True if there are changes not reflected in memory, False otherwise.
-        """
+    def _reload_if_necessary(self):
+        # type: () -> None
+        """Reload cache from persistence layer, if necessary"""
         try:
-            updated = os.path.getmtime(self._cache_location)
-            return self._last_sync < updated
+            if self._last_sync < self._persistence.time_last_modified():
+                self.deserialize(self._persistence.load())
+                self._last_sync = time.time()
         except IOError as exp:
             if exp.errno != errno.ENOENT:
-                raise exp
-            return False
-
-    def _write(self, contents):
-        # type: (str) -> None
-        """Handles actually committing the serialized form of this TokenCache to persisted storage.
-        For types derived of this, class that will be a file, which has the ability to track a last
-        modified time.
-
-        :param contents: The serialized contents of a TokenCache
-        """
-        with open(self._cache_location, 'w+') as handle:
-            handle.write(contents)
-
-    def _read(self):
-        # type: () -> str
-        """Fetches the contents of a file and invokes deserialization."""
-        with open(self._cache_location, 'r') as handle:
-            return handle.read()
+                raise
+            # Otherwise, from cache's perspective, a nonexistent file is a NO-OP
 
     def modify(self, credential_type, old_entry, new_key_value_pairs=None):
         with CrossPlatLock(self._lock_location):
-            if self._needs_refresh():
-                try:
-                    self.deserialize(self._read())
-                except IOError as exp:
-                    if exp.errno != errno.ENOENT:
-                        raise
-            super(FileTokenCache, self).modify(
+            self._reload_if_necessary()
+            super(PersistedTokenCache, self).modify(
                 credential_type,
                 old_entry,
                 new_key_value_pairs=new_key_value_pairs)
-            self._write(self.serialize())
-            self._last_sync = os.path.getmtime(self._cache_location)
+            self._persistence.save(self.serialize())
+            self._last_sync = time.time()
 
     def find(self, credential_type, **kwargs):  # pylint: disable=arguments-differ
         with CrossPlatLock(self._lock_location):
-            if self._needs_refresh():
-                try:
-                    self.deserialize(self._read())
-                except IOError as exp:
-                    if exp.errno != errno.ENOENT:
-                        raise
-                self._last_sync = time.time()
-            return super(FileTokenCache, self).find(credential_type, **kwargs)
+            self._reload_if_necessary()
+            return super(PersistedTokenCache, self).find(credential_type, **kwargs)
 
 
-class UnencryptedTokenCache(FileTokenCache):
-    """An unprotected token cache to default to when no-platform specific option is available."""
-    def __init__(self, cache_location, **kwargs):
-        warnings.warn("You are using an unprotected token cache, "
-                      "because an encrypted option is not available for {}".format(sys.platform),
-                      RuntimeWarning)
-        super(UnencryptedTokenCache, self).__init__(cache_location, **kwargs)
+class FileTokenCache(PersistedTokenCache):
+    """A token cache which uses plain text file to store your tokens."""
+    def __init__(self, cache_location, **ignored):  # pylint: disable=unused-argument
+        warnings.warn("You are using an unprotected token cache", RuntimeWarning)
+        warnings.warn("Use PersistedTokenCache(...) instead", DeprecationWarning)
+        super(FileTokenCache, self).__init__(FilePersistence(cache_location))
 
+UnencryptedTokenCache = FileTokenCache  # For backward compatibility
 
-class WindowsTokenCache(FileTokenCache):
-    """A SerializableTokenCache implementation which uses Win32 encryption APIs to protect your
-    tokens.
-    """
-    def __init__(self, cache_location, entropy='', **kwargs):
-        super(WindowsTokenCache, self).__init__(cache_location, **kwargs)
-        self._dp_agent = WindowsDataProtectionAgent(entropy=entropy)
 
-    def _write(self, contents):
-        with open(self._cache_location, 'wb') as handle:
-            handle.write(self._dp_agent.protect(contents))
+class WindowsTokenCache(PersistedTokenCache):
+    """A token cache which uses Windows DPAPI to encrypt your tokens."""
+    def __init__(
+            self, cache_location, entropy='',
+            **ignored):  # pylint: disable=unused-argument
+        warnings.warn("Use PersistedTokenCache(...) instead", DeprecationWarning)
+        super(WindowsTokenCache, self).__init__(
+            FilePersistenceWithDataProtection(cache_location, entropy=entropy))
 
-    def _read(self):
-        with open(self._cache_location, 'rb') as handle:
-            cipher_text = handle.read()
-        return self._dp_agent.unprotect(cipher_text)
-
-
-class OSXTokenCache(FileTokenCache):
-    """A SerializableTokenCache implementation which uses native Keychain libraries to protect your
-    tokens.
-    """
 
+class OSXTokenCache(PersistedTokenCache):
+    """A token cache which uses native Keychain libraries to encrypt your tokens."""
     def __init__(self,
                  cache_location,
                  service_name='Microsoft.Developer.IdentityService',
                  account_name='MSALCache',
-                 **kwargs):
-        super(OSXTokenCache, self).__init__(cache_location, **kwargs)
-        self._service_name = service_name
-        self._account_name = account_name
-
-    def _read(self):
-        with Keychain() as locker:
-            return locker.get_generic_password(self._service_name, self._account_name)
-
-    def _write(self, contents):
-        with Keychain() as locker:
-            locker.set_generic_password(self._service_name, self._account_name, contents)
-            with open(self._cache_location, "w+") as handle:
-                handle.write('{} {}'.format(os.getpid(), sys.argv[0]))
+                 **ignored):  # pylint: disable=unused-argument
+        warnings.warn("Use PersistedTokenCache(...) instead", DeprecationWarning)
+        super(OSXTokenCache, self).__init__(
+            KeychainPersistence(cache_location, service_name, account_name))
+
diff --git a/sample/token_cache_sample.py b/sample/token_cache_sample.py
@@ -0,0 +1,39 @@
+import sys
+import logging
+import json
+
+from msal_extensions import *
+
+
+def build_persistence(location, fallback_to_plaintext=False):
+    """Build a suitable persistence instance based your current OS"""
+    if sys.platform.startswith('win'):
+        return FilePersistenceWithDataProtection(location)
+    if sys.platform.startswith('darwin'):
+        return KeychainPersistence(location, "my_service_name", "my_account_name")
+    if sys.platform.startswith('linux'):
+        try:
+            return LibsecretPersistence(
+                # By using same location as the fall back option below,
+                # this would override the unencrypted data stored by the
+                # fall back option.  It is probably OK, or even desirable
+                # (in order to aggressively wipe out plain-text persisted data),
+                # unless there would frequently be a desktop session and
+                # a remote ssh session being active simultaneously.
+                location,
+                schema_name="my_schema_name",
+                attributes={"my_attr1": "foo", "my_attr2": "bar"},
+                )
+        except:  # pylint: disable=bare-except
+            if not fallback_to_plaintext:
+                raise
+            logging.exception("Encryption unavailable. Opting in to plain text.")
+    return FilePersistence(location)
+
+persistence = build_persistence("token_cache.bin")
+print("Is this persistence encrypted?", persistence.is_encrypted)
+
+cache = PersistedTokenCache(persistence)
+# Now you can use it in an msal application like this:
+#   app = msal.PublicClientApplication("my_client_id", token_cache=cache)
+
diff --git a/tests/test_agnostic_backend.py b/tests/test_agnostic_backend.py
@@ -1,54 +1,46 @@
 import os
 import shutil
 import tempfile
-import pytest
+import sys
+
 import msal
+import pytest
 
+from msal_extensions import *
 
-def test_file_token_cache_roundtrip():
-    from msal_extensions.token_cache import FileTokenCache
 
+@pytest.fixture
+def temp_location():
+    test_folder = tempfile.mkdtemp(prefix="test_token_cache_roundtrip")
+    yield os.path.join(test_folder, 'token_cache.bin')
+    shutil.rmtree(test_folder, ignore_errors=True)
+
+
+def _test_token_cache_roundtrip(cache):
     client_id = os.getenv('AZURE_CLIENT_ID')
     client_secret = os.getenv('AZURE_CLIENT_SECRET')
     if not (client_id and client_secret):
-        pytest.skip('no credentials present to test FileTokenCache round-trip with.')
-
-    test_folder = tempfile.mkdtemp(prefix="msal_extension_test_file_token_cache_roundtrip")
-    cache_file = os.path.join(test_folder, 'msal.cache')
-    try:
-        subject = FileTokenCache(cache_location=cache_file)
-        app = msal.ConfidentialClientApplication(
-            client_id=client_id,
-            client_credential=client_secret,
-            token_cache=subject)
-        desired_scopes = ['https://graph.microsoft.com/.default']
-        token1 = app.acquire_token_for_client(scopes=desired_scopes)
-        os.utime(cache_file, None)  # Mock having another process update the cache.
-        token2 = app.acquire_token_silent(scopes=desired_scopes, account=None)
-        assert token1['access_token'] == token2['access_token']
-    finally:
-        shutil.rmtree(test_folder, ignore_errors=True)
-
-
-def test_current_platform_cache_roundtrip():
+        pytest.skip('no credentials present to test TokenCache round-trip with.')
+
+    app = msal.ConfidentialClientApplication(
+        client_id=client_id,
+        client_credential=client_secret,
+        token_cache=cache)
+    desired_scopes = ['https://graph.microsoft.com/.default']
+    token1 = app.acquire_token_for_client(scopes=desired_scopes)
+    os.utime(  # Mock having another process update the cache
+        cache._persistence.get_location(), None)
+    token2 = app.acquire_token_silent(scopes=desired_scopes, account=None)
+    assert token1['access_token'] == token2['access_token']
+
+def test_file_token_cache_roundtrip(temp_location):
+    from msal_extensions.token_cache import FileTokenCache
+    _test_token_cache_roundtrip(FileTokenCache(temp_location))
+
+def test_current_platform_cache_roundtrip_with_alias_class(temp_location):
     from msal_extensions import TokenCache
-    client_id = os.getenv('AZURE_CLIENT_ID')
-    client_secret = os.getenv('AZURE_CLIENT_SECRET')
-    if not (client_id and client_secret):
-        pytest.skip('no credentials present to test FileTokenCache round-trip with.')
-
-    test_folder = tempfile.mkdtemp(prefix="msal_extension_test_file_token_cache_roundtrip")
-    cache_file = os.path.join(test_folder, 'msal.cache')
-    try:
-        subject = TokenCache(cache_location=cache_file)
-        app = msal.ConfidentialClientApplication(
-            client_id=client_id,
-            client_credential=client_secret,
-            token_cache=subject)
-        desired_scopes = ['https://graph.microsoft.com/.default']
-        token1 = app.acquire_token_for_client(scopes=desired_scopes)
-        os.utime(cache_file, None)  # Mock having another process update the cache.
-        token2 = app.acquire_token_silent(scopes=desired_scopes, account=None)
-        assert token1['access_token'] == token2['access_token']
-    finally:
-        shutil.rmtree(test_folder, ignore_errors=True)
+    _test_token_cache_roundtrip(TokenCache(temp_location))
+
+def test_persisted_token_cache(temp_location):
+    _test_token_cache_roundtrip(PersistedTokenCache(FilePersistence(temp_location)))
+
