components/ui-request-interceptor/io.entgra.device.mgt.core.ui.request.interceptor/pom.xml

@@ -319,5 +319,20 @@
             <artifactId>httpmime</artifactId>
             <scope>compile</scope>
         </dependency>
+        <dependency>
+            <groupId>org.wso2.carbon.identity.framework</groupId>
+            <artifactId>org.wso2.carbon.user.mgt</artifactId>
+            <scope>provided</scope>
+            <exclusions>
+                <exclusion>
+                    <groupId>org.slf4j</groupId>
+                    <artifactId>slf4j-api</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>org.slf4j</groupId>
+                    <artifactId>jcl-over-slf4j</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
     </dependencies>
 </project>

components/ui-request-interceptor/io.entgra.device.mgt.core.ui.request.interceptor/src/main/java/io/entgra/device/mgt/core/ui/request/interceptor/SsoLoginCallbackHandler.java

@@ -21,7 +21,10 @@ package io.entgra.device.mgt.core.ui.request.interceptor;
 import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
+import com.nimbusds.jwt.JWT;
+import com.nimbusds.jwt.JWTParser;
 import io.entgra.device.mgt.core.ui.request.interceptor.beans.AuthData;
+import io.entgra.device.mgt.core.ui.request.interceptor.beans.ProxyResponse;
 import io.entgra.device.mgt.core.ui.request.interceptor.util.HandlerConstants;
 import io.entgra.device.mgt.core.ui.request.interceptor.util.HandlerUtil;
 import org.apache.commons.logging.Log;
@@ -31,15 +34,22 @@ import org.apache.http.HttpStatus;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
-import io.entgra.device.mgt.core.ui.request.interceptor.beans.ProxyResponse;
+import org.wso2.carbon.context.CarbonContext;
+import org.wso2.carbon.context.PrivilegedCarbonContext;
+import org.wso2.carbon.user.api.UserRealm;
+import org.wso2.carbon.user.api.UserStoreException;
+import org.wso2.carbon.user.api.UserStoreManager;
+import org.wso2.carbon.user.core.service.RealmService;
+import org.wso2.carbon.user.mgt.UserRealmProxy;
+import org.wso2.carbon.user.mgt.common.UIPermissionNode;
+import org.wso2.carbon.user.mgt.common.UserAdminException;
 
 import javax.servlet.annotation.MultipartConfig;
 import javax.servlet.annotation.WebServlet;
-import javax.servlet.http.HttpServlet;
+import javax.servlet.http.*;
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
-import javax.servlet.http.HttpSession;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 
 @MultipartConfig
@@ -93,6 +103,8 @@ public class SsoLoginCallbackHandler extends HttpServlet {
         JsonElement jTokenResult = jsonParser.parse(tokenResultResponse.getData());
         if (jTokenResult.isJsonObject()) {
             JsonObject jTokenResultAsJsonObject = jTokenResult.getAsJsonObject();
+            String id_token = jTokenResultAsJsonObject.get("id_token").getAsString();
+            String username = getUsernameFromJwt(id_token);
             AuthData authData = new AuthData();
             authData.setClientId(session.getAttribute("clientId").toString());
             authData.setClientSecret(session.getAttribute("clientSecret").toString());
@@ -101,7 +113,130 @@ public class SsoLoginCallbackHandler extends HttpServlet {
             authData.setRefreshToken(jTokenResultAsJsonObject.get("refresh_token").getAsString());
             authData.setScope(jTokenResultAsJsonObject.get("scope").getAsString());
             session.setAttribute(HandlerConstants.SESSION_AUTH_DATA_KEY, authData);
+            try {
+                if (isUserCanLogin(username)) {
                     resp.sendRedirect(session.getAttribute("redirectUrl").toString());
+                } else {
+                    for (String path : HandlerConstants.SSO_LOGOUT_COOKIE_PATHS) {
+                        removeCookie(HandlerConstants.JSESSIONID_KEY, path, resp);
+                    }
+                    removeCookie(HandlerConstants.COMMON_AUTH_ID_KEY, "/", resp);
+                    ProxyResponse proxyResponse = new ProxyResponse();
+                    proxyResponse.setStatus(ProxyResponse.Status.SUCCESS);
+                    proxyResponse.setCode(HttpStatus.SC_OK);
+
+                    if (session != null) {
+                        session.invalidate();
+                    }
+                    resp.sendRedirect("/authenticationendpoint/oauth2_error.do");
+                }
+            } catch (UserStoreException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    public static String getUsernameFromJwt(String jwtString) {
+        try {
+            // Parse the JWT string
+            JWT jwt = JWTParser.parse(jwtString);
+            return jwt.getJWTClaimsSet().getSubject();
+        } catch (Exception e) {
+            return null;
+        }
+    }
+
+    private boolean isUserCanLogin(String username) throws UserStoreException {
+
+        UserStoreManager userStoreManager;
+        RealmService realmService;
+        PrivilegedCarbonContext ctx = PrivilegedCarbonContext.getThreadLocalCarbonContext();
+        UserRealm realm;
+        realmService = (RealmService) ctx.getOSGiService(RealmService.class, null);
+
+        if (realmService == null) {
+            throw new IllegalStateException("Realm service not initialized");
+        }
+        int tenantId = CarbonContext.getThreadLocalCarbonContext().getTenantId();
+        realm = realmService.getTenantUserRealm(tenantId);
+
+        userStoreManager = realmService.getTenantUserRealm(tenantId).getUserStoreManager();
+
+
+        if (!userStoreManager.isExistingUser(username)) {
+            String message = "User by username: " + username + " does not exist for permission retrieval.";
+            log.error(message);
+        }
+        // Get a list of roles which the user assigned to
+        List<String> roles = getFilteredRoles(userStoreManager, username);
+        List<String> permissions = new ArrayList<>();
+        tenantId = PrivilegedCarbonContext.getThreadLocalCarbonContext().getTenantId();
+        // Get permissions for each role
+        for (String roleName : roles) {
+            try {
+                permissions.addAll(getPermissionsListFromRole(roleName, realm, tenantId));
+            } catch (UserAdminException e) {
+                String message = "Error occurred while retrieving the permissions of role '" + roleName + "'";
+                log.error(message, e);
+            }
+        }
+        String targetPermission = "/permission/admin/login";
+        return permissions.contains(targetPermission);
+    }
+
+    private static void removeCookie(String cookieName, String path, HttpServletResponse response) {
+        Cookie cookie = new Cookie(cookieName, "");
+        cookie.setPath(path);
+        cookie.setValue(null);
+        cookie.setMaxAge(0);
+        response.addCookie(cookie);
+    }
+
+    private List<String> getFilteredRoles(UserStoreManager userStoreManager, String username)
+            throws UserStoreException {
+        String[] roleListOfUser;
+        roleListOfUser = userStoreManager.getRoleListOfUser(username);
+        List<String> filteredRoles = new ArrayList<>();
+        for (String role : roleListOfUser) {
+            if (!(role.startsWith("Internal/") || role.startsWith("Authentication/"))) {
+                filteredRoles.add(role);
+            }
+        }
+        return filteredRoles;
+    }
+
+    private static List<String> getPermissionsListFromRole(String roleName, UserRealm userRealm, int tenantId)
+            throws UserAdminException {
+        org.wso2.carbon.user.core.UserRealm userRealmCore;
+        try {
+            userRealmCore = (org.wso2.carbon.user.core.UserRealm) userRealm;
+        } catch (ClassCastException e) {
+            String message = "Provided UserRealm object is not an instance of org.wso2.carbon.user.core.UserRealm";
+            log.error(message, e);
+            throw new UserAdminException(message, e);
+        }
+        UserRealmProxy userRealmProxy = new UserRealmProxy(userRealmCore);
+        List<String> permissionsList = new ArrayList<>();
+        final UIPermissionNode rolePermissions = userRealmProxy.getRolePermissions(roleName, tenantId);
+        iteratePermissions(rolePermissions, permissionsList);
+        return permissionsList;
+    }
+
+    public static void iteratePermissions(UIPermissionNode uiPermissionNode, List<String> list) {
+        // To prevent NullPointer exceptions
+        if (uiPermissionNode == null) {
+            return;
+        }
+        for (UIPermissionNode permissionNode : uiPermissionNode.getNodeList()) {
+            if (permissionNode != null) {
+                if (permissionNode.isSelected()) {
+                    list.add(permissionNode.getResourcePath());
+                }
+                if (permissionNode.getNodeList() != null
+                        && permissionNode.getNodeList().length > 0) {
+                    iteratePermissions(permissionNode, list);
+                }
+            }
         }
     }
 }