Skip to content

Commit

Permalink
GEN-927 - Add bot default roles
Browse files Browse the repository at this point in the history
  • Loading branch information
pmbrull committed Oct 14, 2024
1 parent 9496da1 commit 967e54d
Show file tree
Hide file tree
Showing 5 changed files with 99 additions and 43 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
@Slf4j
public class RoleRepository extends EntityRepository<Role> {
public static final String DOMAIN_ONLY_ACCESS_ROLE = "DomainOnlyAccessRole";
public static final String DEFAULT_BOT_ROLE = "DefaultBotRole";

public RoleRepository() {
super(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,15 @@
import static javax.ws.rs.core.Response.Status.FORBIDDEN;
import static javax.ws.rs.core.Response.Status.OK;
import static org.openmetadata.common.utils.CommonUtil.listOf;
import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
import static org.openmetadata.schema.api.teams.CreateUser.CreatePasswordType.ADMIN_CREATE;
import static org.openmetadata.schema.auth.ChangePasswordRequest.RequestType.SELF;
import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType.BASIC;
import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType.JWT;
import static org.openmetadata.schema.type.Include.ALL;
import static org.openmetadata.service.exception.CatalogExceptionMessage.EMAIL_SENDING_ISSUE;
import static org.openmetadata.service.jdbi3.RoleRepository.DEFAULT_BOT_ROLE;
import static org.openmetadata.service.jdbi3.RoleRepository.DOMAIN_ONLY_ACCESS_ROLE;
import static org.openmetadata.service.jdbi3.UserRepository.AUTH_MECHANISM_FIELD;
import static org.openmetadata.service.secrets.ExternalSecretsManager.NULL_SECRET_STRING;
import static org.openmetadata.service.security.jwt.JWTTokenGenerator.getExpiryDate;
Expand Down Expand Up @@ -51,6 +54,7 @@
import java.io.IOException;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.util.ArrayList;
import java.util.Base64;
import java.util.Date;
import java.util.List;
Expand Down Expand Up @@ -104,7 +108,6 @@
import org.openmetadata.schema.auth.RegistrationRequest;
import org.openmetadata.schema.auth.RevokePersonalTokenRequest;
import org.openmetadata.schema.auth.RevokeTokenRequest;
import org.openmetadata.schema.auth.SSOAuthMechanism;
import org.openmetadata.schema.auth.ServiceTokenType;
import org.openmetadata.schema.auth.TokenRefreshRequest;
import org.openmetadata.schema.auth.TokenType;
Expand All @@ -125,6 +128,7 @@
import org.openmetadata.service.exception.EntityNotFoundException;
import org.openmetadata.service.jdbi3.CollectionDAO;
import org.openmetadata.service.jdbi3.ListFilter;
import org.openmetadata.service.jdbi3.RoleRepository;
import org.openmetadata.service.jdbi3.TokenRepository;
import org.openmetadata.service.jdbi3.UserRepository;
import org.openmetadata.service.jdbi3.UserRepository.UserCsv;
Expand Down Expand Up @@ -173,6 +177,7 @@ public class UserResource extends EntityResource<User, UserRepository> {
public static final String USER_PROTECTED_FIELDS = "authenticationMechanism";
private final JWTTokenGenerator jwtTokenGenerator;
private final TokenRepository tokenRepository;
private final RoleRepository roleRepository;
private AuthenticationConfiguration authenticationConfiguration;
private AuthorizerConfiguration authorizerConfiguration;
private final AuthenticatorHandler authHandler;
Expand All @@ -197,6 +202,7 @@ public UserResource(
jwtTokenGenerator = JWTTokenGenerator.getInstance();
allowedFields.remove(USER_PROTECTED_FIELDS);
tokenRepository = Entity.getTokenRepository();
roleRepository = Entity.getRoleRepository();
UserTokenCache.initialize();
authHandler = authenticatorHandler;
}
Expand Down Expand Up @@ -567,6 +573,7 @@ public Response createUser(
User user = getUser(securityContext.getUserPrincipal().getName(), create);
if (Boolean.TRUE.equals(user.getIsBot())) {
addAuthMechanismToBot(user, create, uriInfo);
addRolesToBot(user, uriInfo);
}

//
Expand Down Expand Up @@ -696,8 +703,8 @@ public Response createOrUpdateUser(
new OperationContext(entityType, EntityUtil.createOrUpdateOperation(resourceContext));
authorizer.authorize(securityContext, createOperationContext, resourceContext);
}
if (Boolean.TRUE.equals(create.getIsBot())) { // TODO expect bot to be created separately
return createOrUpdateBot(user, create, uriInfo, securityContext);
if (Boolean.TRUE.equals(create.getIsBot())) {
return createOrUpdateBotUser(user, create, uriInfo, securityContext);
}
PutResponse<User> response = repository.createOrUpdate(uriInfo, user);
addHref(uriInfo, response.getEntity());
Expand Down Expand Up @@ -1454,7 +1461,7 @@ public void validateEmailAlreadyExists(String email) {
}
}

private Response createOrUpdateBot(
private Response createOrUpdateBotUser(
User user, CreateUser create, UriInfo uriInfo, SecurityContext securityContext) {
User original = retrieveBotUser(user, uriInfo);
String botName = create.getBotName();
Expand Down Expand Up @@ -1488,8 +1495,9 @@ && userHasRelationshipWithAnyBot(original, bot)) {
original.getAuthenticationMechanism());
user.setRoles(original.getRoles());
}
// TODO remove this
// TODO remove this -> Still valid TODO?
addAuthMechanismToBot(user, create, uriInfo);
addRolesToBot(user, uriInfo);
PutResponse<User> response = repository.createOrUpdate(uriInfo, user);
decryptOrNullify(securityContext, response.getEntity());
return response.toResponse();
Expand Down Expand Up @@ -1531,41 +1539,54 @@ private List<CollectionDAO.EntityRelationshipRecord> retrieveBotRelationshipsFor
return repository.findToRecords(bot.getId(), Entity.BOT, Relationship.CONTAINS, Entity.USER);
}

// TODO remove this
// TODO remove this -> still valid TODO?
private void addAuthMechanismToBot(User user, @Valid CreateUser create, UriInfo uriInfo) {
if (!Boolean.TRUE.equals(user.getIsBot())) {
throw new IllegalArgumentException(
"Authentication mechanism change is only supported for bot users");
}
if (isValidAuthenticationMechanism(create)) {
AuthenticationMechanism authMechanism = create.getAuthenticationMechanism();
AuthenticationMechanism.AuthType authType = authMechanism.getAuthType();
switch (authType) {
case JWT -> {
User original = retrieveBotUser(user, uriInfo);
if (original == null
|| !hasAJWTAuthMechanism(user, original.getAuthenticationMechanism())) {
JWTAuthMechanism jwtAuthMechanism =
JsonUtils.convertValue(authMechanism.getConfig(), JWTAuthMechanism.class);
authMechanism.setConfig(
jwtTokenGenerator.generateJWTToken(user, jwtAuthMechanism.getJWTTokenExpiry()));
} else {
authMechanism = original.getAuthenticationMechanism();
}
}
case SSO -> {
SSOAuthMechanism ssoAuthMechanism =
JsonUtils.convertValue(authMechanism.getConfig(), SSOAuthMechanism.class);
authMechanism.setConfig(ssoAuthMechanism);
}
default -> throw new IllegalArgumentException(
String.format("Not supported authentication mechanism type: [%s]", authType.value()));
}
user.setAuthenticationMechanism(authMechanism);
AuthenticationMechanism authMechanism = create.getAuthenticationMechanism();
User original = retrieveBotUser(user, uriInfo);
if (original == null || !hasAJWTAuthMechanism(user, original.getAuthenticationMechanism())) {
JWTAuthMechanism jwtAuthMechanism =
JsonUtils.convertValue(authMechanism.getConfig(), JWTAuthMechanism.class);
authMechanism.setConfig(
jwtTokenGenerator.generateJWTToken(user, jwtAuthMechanism.getJWTTokenExpiry()));
} else {
throw new IllegalArgumentException(
String.format("Authentication mechanism is empty bot user: [%s]", user.getName()));
authMechanism = original.getAuthenticationMechanism();
}
user.setAuthenticationMechanism(authMechanism);
}

private void addRolesToBot(User user, UriInfo uriInfo) {
if (!Boolean.TRUE.equals(user.getIsBot())) {
throw new IllegalArgumentException("Bot roles are only supported for bot users");
}
User original = retrieveBotUser(user, uriInfo);
ArrayList<EntityReference> defaultBotRoles = getDefaultBotRoles(user);
// Keep the incoming roles of the created user
if (!nullOrEmpty(user.getRoles())) {
defaultBotRoles.addAll(user.getRoles());
}
// If user existed, merge roles
if (original != null && !nullOrEmpty(original.getRoles())) {
defaultBotRoles.addAll(original.getRoles());
}
user.setRoles(defaultBotRoles);
}

private ArrayList<EntityReference> getDefaultBotRoles(User user) {
ArrayList<EntityReference> defaultBotRoles = new ArrayList<>();
EntityReference defaultBotRole =
roleRepository.getReferenceByName(DEFAULT_BOT_ROLE, Include.NON_DELETED);
defaultBotRoles.add(defaultBotRole);

if (!nullOrEmpty(user.getDomains())) {
EntityReference domainOnlyAccessRole =
roleRepository.getReferenceByName(DOMAIN_ONLY_ACCESS_ROLE, Include.NON_DELETED);
defaultBotRoles.add(domainOnlyAccessRole);
}
return defaultBotRoles;
}

@Nullable
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
{
"name": "DefaultBotRole",
"displayName": "Default Bot Role",
"description": "Role Corresponding to a Bot by default.",
"allowDelete": false,
"provider": "system",
"policies" : [
{
"type" : "policy",
"name" : "DefaultBotPolicy"
},
{
"type" : "policy",
"name" : "DataConsumerPolicy"
}
]
}
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,8 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
public static EntityReference USER2_REF;
public static User USER_TEAM21;
public static User BOT_USER;
public static EntityReference DEFAULT_BOT_ROLE_REF;
public static EntityReference DOMAIN_ONLY_ACCESS_ROLE_REF;

public static Team ORG_TEAM;
public static Team TEAM1;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@
import static org.openmetadata.service.exception.CatalogExceptionMessage.notAdmin;
import static org.openmetadata.service.exception.CatalogExceptionMessage.operationNotAllowed;
import static org.openmetadata.service.exception.CatalogExceptionMessage.permissionNotAllowed;
import static org.openmetadata.service.jdbi3.RoleRepository.DEFAULT_BOT_ROLE;
import static org.openmetadata.service.jdbi3.RoleRepository.DOMAIN_ONLY_ACCESS_ROLE;
import static org.openmetadata.service.resources.teams.UserResource.USER_PROTECTED_FIELDS;
import static org.openmetadata.service.security.SecurityUtil.authHeaders;
import static org.openmetadata.service.util.EntityUtil.fieldAdded;
Expand Down Expand Up @@ -109,18 +111,17 @@
import org.openmetadata.schema.auth.RegistrationRequest;
import org.openmetadata.schema.auth.RevokePersonalTokenRequest;
import org.openmetadata.schema.auth.RevokeTokenRequest;
import org.openmetadata.schema.auth.SSOAuthMechanism;
import org.openmetadata.schema.entity.data.Table;
import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
import org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType;
import org.openmetadata.schema.entity.teams.Role;
import org.openmetadata.schema.entity.teams.Team;
import org.openmetadata.schema.entity.teams.User;
import org.openmetadata.schema.security.client.GoogleSSOClientConfig;
import org.openmetadata.schema.type.ApiStatus;
import org.openmetadata.schema.type.ChangeDescription;
import org.openmetadata.schema.type.EntityReference;
import org.openmetadata.schema.type.ImageList;
import org.openmetadata.schema.type.Include;
import org.openmetadata.schema.type.MetadataOperation;
import org.openmetadata.schema.type.Profile;
import org.openmetadata.schema.type.Webhook;
Expand All @@ -129,6 +130,7 @@
import org.openmetadata.service.Entity;
import org.openmetadata.service.auth.JwtResponse;
import org.openmetadata.service.exception.CatalogExceptionMessage;
import org.openmetadata.service.jdbi3.RoleRepository;
import org.openmetadata.service.jdbi3.TeamRepository.TeamCsv;
import org.openmetadata.service.jdbi3.UserRepository.UserCsv;
import org.openmetadata.service.resources.EntityResourceTest;
Expand All @@ -150,11 +152,13 @@ public class UserResourceTest extends EntityResourceTest<User, CreateUser> {
private static final Profile PROFILE =
new Profile().withImages(new ImageList().withImage(URI.create("https://image.com")));
private static final TeamResourceTest TEAM_TEST = new TeamResourceTest();
private final RoleRepository roleRepository;

public UserResourceTest() {
super(USER, User.class, UserList.class, "users", UserResource.FIELDS);
supportedNameCharacters = "_-.";
supportsSearchIndex = true;
roleRepository = Entity.getRoleRepository();
}

public void setupUsers(TestInfo test) throws HttpResponseException {
Expand Down Expand Up @@ -193,6 +197,11 @@ public void setupUsers(TestInfo test) throws HttpResponseException {
Set<String> userFields = Entity.getEntityFields(User.class);
userFields.remove("authenticationMechanism");
BOT_USER = getEntityByName(INGESTION_BOT, String.join(",", userFields), ADMIN_AUTH_HEADERS);

// Get the bot roles
DEFAULT_BOT_ROLE_REF = roleRepository.getReferenceByName(DEFAULT_BOT_ROLE, Include.NON_DELETED);
DOMAIN_ONLY_ACCESS_ROLE_REF =
roleRepository.getReferenceByName(DOMAIN_ONLY_ACCESS_ROLE, Include.NON_DELETED);
}

@Test
Expand Down Expand Up @@ -886,28 +895,26 @@ protected void validateCommonEntityFields(User entity, CreateEntity create, Stri
void put_generateToken_bot_user_200_ok() throws HttpResponseException {
AuthenticationMechanism authMechanism =
new AuthenticationMechanism()
.withAuthType(AuthType.SSO)
.withConfig(
new SSOAuthMechanism()
.withSsoServiceType(SSOAuthMechanism.SsoServiceType.GOOGLE)
.withAuthConfig(
new GoogleSSOClientConfig().withSecretKey("/path/to/secret.json")));
.withAuthType(AuthType.JWT)
.withConfig(new JWTAuthMechanism().withJWTTokenExpiry(JWTTokenExpiry.Unlimited));
CreateUser create =
createBotUserRequest("ingestion-bot-jwt")
.withEmail("[email protected]")
.withRoles(List.of(ROLE1_REF.getId()))
.withAuthenticationMechanism(authMechanism);
User user = createEntity(create, USER_WITH_CREATE_HEADERS);
user = getEntity(user.getId(), "*", ADMIN_AUTH_HEADERS);
assertEquals(1, user.getRoles().size());
// Has the given role and the default bot role
assertEquals(2, user.getRoles().size());
TestUtils.put(
getResource(String.format("users/generateToken/%s", user.getId())),
new GenerateTokenRequest().withJWTTokenExpiry(JWTTokenExpiry.Seven),
OK,
ADMIN_AUTH_HEADERS);
user = getEntity(user.getId(), "*", ADMIN_AUTH_HEADERS);
assertNull(user.getAuthenticationMechanism());
assertEquals(1, user.getRoles().size());
// Has the given role and the default bot role
assertEquals(2, user.getRoles().size());
JWTAuthMechanism jwtAuthMechanism =
TestUtils.get(
getResource(String.format("users/token/%s", user.getId())),
Expand Down Expand Up @@ -1441,6 +1448,14 @@ public void validateCreatedEntity(
for (UUID roleId : listOrEmpty(createRequest.getRoles())) {
expectedRoles.add(new EntityReference().withId(roleId).withType(Entity.ROLE));
}

// bots are created with default roles
if (createRequest.getIsBot()) {
expectedRoles.add(DEFAULT_BOT_ROLE_REF);
if (!nullOrEmpty(createRequest.getDomains())) {
expectedRoles.add(DOMAIN_ONLY_ACCESS_ROLE_REF);
}
}
assertRoles(user, expectedRoles);

List<EntityReference> expectedTeams = new ArrayList<>();
Expand Down

0 comments on commit 967e54d

Please sign in to comment.