% CredentialStore - Store for secure user credentials.
% A user name is associated with at most one credential. If a credential is
% added for an existing user name, the existing credential is replaced even
% if the lifetime is different.
% Credentials can be stored for three different lifetimes:
% Temporary - Credential is removed after a fixed duration
% Session - Credential persists as long as current MATLAB session
% Always - Credential is saved to disk and reloaded each session
% All credentials are unit8 arrays. This provides maximum version
% compatibility between clients.

% Copyright 2019-2021 The MathWorks, Inc.

classdef CredentialStore < handle

    properties (GetAccess = private, Constant)
        CredentialDirectoryName = 'jobmanager_credentials';
    end
    
    properties (SetAccess = private, GetAccess = private)
        % CredentialMap - Map of user name to credential
        CredentialMap
        % ExpirationDateMap - Map of user name to credential expiration datetimes
        ExpirationDateMap
        % PersistentCredentialMap - Subset of CredentialMap for credentials
        % which are persisted to disk in StoreFile
        PersistentCredentialMap
    end
    
    properties (SetAccess = immutable)
        % StoreFile - File where credential store is persisted to disk
        StoreFile
        % Lifetime - Duration after which expirable credentials expire
        Lifetime
    end
   
    methods (Hidden)
        % Constructor with specific file used by testing
        function obj = CredentialStore(storeFile, lifetime)
            obj.PersistentCredentialMap = containers.Map('KeyType', 'char', 'ValueType', 'any');
            obj.StoreFile = storeFile;
            obj.restore();
            
            obj.Lifetime = lifetime;
            if isempty(obj.PersistentCredentialMap)
                obj.CredentialMap = containers.Map('KeyType', 'char', 'ValueType', 'any');
            else
                % Copy persistent credentials into memory
                obj.CredentialMap = containers.Map(keys(obj.PersistentCredentialMap), ...
                    values(obj.PersistentCredentialMap));
            end
            obj.ExpirationDateMap = containers.Map('KeyType', 'char', 'ValueType', 'any');
        end
    end
    
    methods (Access = private)
        % Save obj.PersistentCredentialMap to disk
        function snapshot(obj)            
            [fid, msg]  = fopen(obj.StoreFile, 'w+');
            if fid == -1
                dctSchedulerMessage(1, ...
                    'Failed to open persistent credentials file %s.\nReason: %s', obj.StoreFile, msg);
                return
            end
            try
                closeFile = onCleanup(@() fclose(fid));
                jsonData = cellfun(@(name) iEncodeAsStruct(name, obj.PersistentCredentialMap(name)), ...
                    keys(obj.PersistentCredentialMap));
                jsonText = jsonencode(jsonData);
                fwrite(fid, jsonText);
            catch E
                dctSchedulerMessage(1, ...
                    'Error occured saving persistent credentials', E);
            end
        end
        
        % Load obj.PersistentCredentialMap from disk
        function restore(obj)
            if ~exist(obj.StoreFile, 'file')
                dctSchedulerMessage(4, ...
                    'Persistent credentials file %s does not exist', obj.StoreFile);
                return
            end
            
            try
                jsonText = fileread(obj.StoreFile);
                if isempty(jsonText)
                    % Nothing written yet
                    return
                end
                jsonData = jsondecode(jsonText);
                for idx = 1:numel(jsonData)
                    data = jsonData(idx);
                    obj.PersistentCredentialMap(data.Name) = uint8(data.Value);
                end
            catch E
                dctSchedulerMessage(1, ...
                    'Error occured loading persistent credentials', E);
            end
        end
        
        % Add a user credential
        function addCredential(obj, username, credential)
            obj.CredentialMap(username) = credential;
        end
        
        % Add a user credential which will expire
        function addExpiringCredential(obj, username, credential)
            expirationDate = datetime('now') + obj.Lifetime;
            obj.addCredential(username, credential);
            obj.ExpirationDateMap(username) = expirationDate;
        end
        
        % Add a user credential which is persisted to disk
        function addPersistentCredential(obj, username, credential)
            % Ensure we are up to date with the file before overwriting
            obj.restore();
            obj.addCredential(username, credential);
            obj.PersistentCredentialMap(username) = credential;
            obj.snapshot();
        end
        
        % Remove all credentials whose expiration time has passed
        function removeExpiredCredentials(obj)
            now = datetime('now');
            expiringUsernames = keys(obj.ExpirationDateMap);
            hasExpired = cellfun(@(name) obj.checkExpiration(name, now), expiringUsernames);
            expiredUsernames = expiringUsernames(hasExpired);
            obj.removeCredentials(expiredUsernames);
        end
        
        % Test if credentials for a given user have expired
        function expired = checkExpiration(obj, username, now)
            expired = false;
            if obj.CredentialMap.isKey(username)
                expirationDate = obj.ExpirationDateMap(username);
                if expirationDate <= now
                    % Credential has expired and must be removed
                    expired = true;
                end
            else
                % Credential is missing so should be considered expired
                expired = true;
            end
        end
        
        % Completely remove credentials for given users. If the
        % credentials were persisted to disk, the file will be updated.
        function removeCredentials(obj, usernames)
            iSafeRemove(obj.ExpirationDateMap, usernames);
            iSafeRemove(obj.CredentialMap, usernames);
            
            updateFile = ~isempty(intersect(usernames, keys(obj.PersistentCredentialMap)));
            if updateFile
                % Ensure we are up to date with the file before overwriting
                obj.restore();
                iSafeRemove(obj.PersistentCredentialMap, usernames);
                obj.snapshot();
            end
        end
        
        % Get a single user credential
        function credential = getCredential(obj, username)
            if obj.CredentialMap.isKey(username)
                credential = obj.CredentialMap(username);
            else
                % Chose to return empty. It's up to the caller to decide if
                % this is an exception
                credential = uint8.empty();
            end
        end
    end
    
    methods
        function credentials = get(obj, usernames)
        %GET Get credentials matching the given user names
        %    CREDENTIALS = GET(OBJ, USERNAMES) returns  credentials stored
        %    for the given USERNAMES. If no credentials exist then an empty
        %    struct is returned.
        %    If USERNAMES is a cell array, CREDENTIALS is a cell
        %    array of credentials

            obj.removeExpiredCredentials();
            isCell = iscell(usernames);
            if ~isCell
                usernames = {usernames};
            end
            credentials = cellfun(@obj.getCredential, usernames, 'UniformOutput', false);
            if ~isCell
                credentials = credentials{1};
            end
        end
        
        function put(obj, username, credential, rememberPassword)
        %PUT Add credential for a given user name
        %    PUT(OBJ, USERNAME, CREDENTIAL, REMEMBERPASSWORD) add the
        %    CREDENTIAL to the store for the USERNAME entry. If
        %    credentials already exist for this user they are replaced.
        %    REMEMBERPASSWORD is an enum indicating how long the credentials
        %    will persist. The default value is Always. Temporary will
        %    cause credentials to only persist as long as the store
        %    lifetime. Session will cause the credential to persist for the
        %    lifetime of the MATLAB session. Always will cause the
        %    credential to be persisted to disk and loaded in future MATLAB
        %    sessions.
        
            import parallel.internal.types.RememberPassword;
            if nargin < 4
                rememberPassword = RememberPassword.Always;
            end
            
            % First remove any credential in case we are switching the
            % lifetime of an existing credential
            obj.removeCredentials(username);
            
            switch rememberPassword
                case RememberPassword.Temporary
                    obj.addExpiringCredential(username, credential);
                case RememberPassword.Session
                    obj.addCredential(username, credential);
                case RememberPassword.Always
                    obj.addPersistentCredential(username, credential);
            end
        end
        
        function remove(obj, usernames)
        %REMOVE Remove any credentials for the given user names
        %   REMOVE(OBJ, USERNAMES) removes any credentials stored for
        %   USERNAMES. If no credentials exists, this method silently returns.
        %   USERNAMES can be either a single user name as a char vector or
        %   string, or multiple user names as a cellstr.
        
            obj.removeCredentials(usernames);
        end
        
        function clear(obj)
        %CLEAR Remove all credentials
        %   CLEAR(OBJ) Remove all credentials from this store. The file
        %   storing credentials to disk will be cleared but not deleted.
        
            obj.CredentialMap = containers.Map('KeyType', 'char', 'ValueType', 'any');
            obj.ExpirationDateMap = containers.Map('KeyType', 'char', 'ValueType', 'any');
            obj.PersistentCredentialMap = containers.Map('KeyType', 'char', 'ValueType', 'any');
            obj.snapshot();
        end
        
        function tf = contains(obj, username)
        %CONTAINS Check if this store contains credentials for a given user
        %name
        %   TF = CONTAINS(OBJ, USERNAME) Returns true if credentials exist
        %   for USERNAME and false otherwise.
        
            obj.removeExpiredCredentials();
            tf = obj.CredentialMap.isKey(username);
        end
        
        function allUsers = getAllUsers(obj)
        %GETALLUSERS Return all user names in this store
        %   ALLUSERS = GETALLUSERS(OBJ) Return all user names in this
        %   store as cellstr.
        
            allUsers = keys(obj.CredentialMap);
        end
    end
    
    methods (Static)
        function store = create(sessionName, lifetime)
        %CREATE Factory method to create a new store
        % STORE = CREATE(SESSIONNAME, LIFETIME) create a new store.
        % SESSIONNAME is used to create the file which persists credentials
        % to disk. If a file matching this session name already exists,
        % those credentials will be loaded into STORE. 
        % LIFETIME is a duration which controls how long any credential
        % with a remember choice of Temporary will persist in this store.
        % Default is 5 minutes
        
            if nargin < 2
                lifetime = minutes(5);
            end
                        
            % Credential directory does not go under prefsdir - if it was 
            % it would be bundled with deployed applications which would 
            % not be right. Instead it goes beside the prefdir. This is
            % the same place local_scheduler_data currently goes.
            rootDir = fileparts(prefdir);
            credDir = fullfile(rootDir, parallel.internal.cluster.CredentialStore.CredentialDirectoryName);
            if exist(credDir, 'dir') ~= 7
                if mkdir(credDir) ~= 1
                    error('parallel:cluster:mjs:MultiUserPersistentCredentialStoreNoCredDir', credDir);
                end
            end
            
            storeFile = sessionName + ".json";
            storeFile = fullfile(credDir, storeFile);
            store = parallel.internal.cluster.CredentialStore(storeFile, lifetime);
        end
    end
end

% Helper method to avoid warning when removing keys which don't exist
function keysToRemove = iSafeRemove(map, keysToRemove)
    allKeys = keys(map);
    keysToRemove = intersect(keysToRemove, allKeys);
    remove(map, keysToRemove);
end

function data = iEncodeAsStruct(name, value)
    data = struct('Name', name, 'Value', value);
end
