content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.ambari.server.security.authentication.jwt;
  
  import static org.easymock.EasyMock.anyObject;
  import static org.easymock.EasyMock.capture;
  import static org.easymock.EasyMock.eq;
  import static org.easymock.EasyMock.expect;
  import static org.easymock.EasyMock.expectLastCall;
  import static org.easymock.EasyMock.newCapture;
  import static org.easymock.EasyMock.verify;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertSame;
  import static org.junit.Assert.assertTrue;
  
  import java.security.KeyPair;
  import java.security.KeyPairGenerator;
  import java.security.NoSuchAlgorithmException;
  import java.security.interfaces.RSAPrivateKey;
  import java.security.interfaces.RSAPublicKey;
  import java.util.Calendar;
  import java.util.Collections;
  import java.util.Date;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  import jakarta.servlet.FilterChain;
  import jakarta.servlet.http.Cookie;
  import jakarta.servlet.http.HttpServletRequest;
  import jakarta.servlet.http.HttpServletResponse;
  
  import org.apache.ambari.server.configuration.AmbariServerConfigurationKey;
  import org.apache.ambari.server.configuration.Configuration;
  import org.apache.ambari.server.orm.entities.UserAuthenticationEntity;
  import org.apache.ambari.server.orm.entities.UserEntity;
  import org.apache.ambari.server.security.AmbariEntryPoint;
  import org.apache.ambari.server.security.authentication.AmbariAuthenticationEventHandler;
  import org.apache.ambari.server.security.authentication.AmbariAuthenticationException;
  import org.apache.ambari.server.security.authentication.AmbariAuthenticationFilter;
  import org.apache.ambari.server.security.authorization.User;
  import org.apache.ambari.server.security.authorization.UserAuthenticationType;
  import org.apache.ambari.server.security.authorization.Users;
  import org.apache.commons.lang.StringUtils;
  import org.easymock.Capture;
  import org.easymock.CaptureType;
  import org.easymock.EasyMockSupport;
  import org.junit.Before;
  import org.junit.BeforeClass;
  import org.junit.Test;
  import org.springframework.security.core.Authentication;
  import org.springframework.security.core.context.SecurityContextHolder;
  import org.springframework.security.web.AuthenticationEntryPoint;
  
  import com.nimbusds.jose.JOSEException;
  import com.nimbusds.jose.JWSAlgorithm;
  import com.nimbusds.jose.JWSHeader;
  import com.nimbusds.jose.crypto.RSASSASigner;
  import com.nimbusds.jwt.JWTClaimsSet;
  import com.nimbusds.jwt.SignedJWT;
  
  public class AmbariJwtAuthenticationFilterTest extends EasyMockSupport {
    private static RSAPublicKey publicKey;
    private static RSAPrivateKey privateKey;
    private static RSAPrivateKey invalidPrivateKey;
  
  
    @BeforeClass
    public static void generateKeyPair() throws NoSuchAlgorithmException {
      KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
      keyPairGenerator.initialize(512);
      KeyPair keyPair = keyPairGenerator.generateKeyPair();
      publicKey = (RSAPublicKey) keyPair.getPublic();
      privateKey = (RSAPrivateKey) keyPair.getPrivate();
  
      keyPair = keyPairGenerator.generateKeyPair();
      invalidPrivateKey = (RSAPrivateKey) keyPair.getPrivate();
    }
  
    @Before
    public void setup() {
      SecurityContextHolder.clearContext();
    }
  
    private JwtAuthenticationProperties createTestProperties() {
      return createTestProperties(Collections.singletonList("test-audience"));
    }
  
    private JwtAuthenticationProperties createTestProperties(List<String> audiences) {
      final Map<String, String> configurationMap = new HashMap<>();
      configurationMap.put(AmbariServerConfigurationKey.SSO_JWT_COOKIE_NAME.key(), "non-default");
      configurationMap.put(AmbariServerConfigurationKey.SSO_JWT_AUDIENCES.key(), audiences == null || audiences.isEmpty() ? "" : StringUtils.join(audiences, ","));
      configurationMap.put(AmbariServerConfigurationKey.SSO_AUTHENTICATION_ENABLED.key(), "true");
      JwtAuthenticationProperties properties = new JwtAuthenticationProperties(configurationMap);
      properties.setPublicKey(publicKey);
  
      return properties;
    }
  
    private SignedJWT getSignedToken() throws JOSEException {
      return getSignedToken("test-audience");
    }
  
    private SignedJWT getSignedToken(String audience) throws JOSEException {
      Calendar calendar = Calendar.getInstance();
      calendar.setTimeInMillis(System.currentTimeMillis());
      calendar.add(Calendar.DATE, 1); //add one day
      return getSignedToken(calendar.getTime(), audience);
    }
  
    private SignedJWT getSignedToken(Date expirationTime, String audience) throws JOSEException {
      RSASSASigner signer = new RSASSASigner(privateKey);
  
      Calendar calendar = Calendar.getInstance();
      calendar.setTimeInMillis(System.currentTimeMillis());
      JWTClaimsSet claimsSet = new JWTClaimsSet.Builder()
          .subject("test-user")
          .issuer("unit-test")
          .issueTime(calendar.getTime())
          .expirationTime(expirationTime)
          .audience(audience)
          .build();
  
      SignedJWT signedJWT = new SignedJWT(new JWSHeader(JWSAlgorithm.RS256), claimsSet);
      signedJWT.sign(signer);
  
      return signedJWT;
    }
  
    private SignedJWT getInvalidToken() throws JOSEException {
      RSASSASigner signer = new RSASSASigner(invalidPrivateKey);
  
      Calendar issueTime = Calendar.getInstance();
      issueTime.setTimeInMillis(System.currentTimeMillis());
      issueTime.add(Calendar.DATE, -2);
  
      Calendar expirationTime = Calendar.getInstance();
      issueTime.setTimeInMillis(System.currentTimeMillis());
      expirationTime.add(Calendar.DATE, -1);
  
      JWTClaimsSet claimsSet = new JWTClaimsSet.Builder()
          .subject("test-user")
          .issuer("unit-test")
          .issueTime(issueTime.getTime())
          .expirationTime(issueTime.getTime())
          .audience("test-audience-invalid")
          .build();
  
      SignedJWT signedJWT = new SignedJWT(new JWSHeader(JWSAlgorithm.RS256), claimsSet);
      signedJWT.sign(signer);
  
      return signedJWT;
    }
  
    @Test
    public void testGetJWTFromCookie() throws Exception {
      HttpServletRequest request = createNiceMock(HttpServletRequest.class);
      Cookie cookie = createNiceMock(Cookie.class);
  
      expect(cookie.getName()).andReturn("non-default");
      expect(cookie.getValue()).andReturn("stubtokenstring");
  
      expect(request.getCookies()).andReturn(new Cookie[]{cookie});
  
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      String jwtFromCookie = filter.getJWTFromCookie(request);
  
      verifyAll();
  
      assertEquals("stubtokenstring", jwtFromCookie);
    }
  
    @Test
    public void testValidateSignature() throws Exception {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertTrue(filter.validateSignature(getSignedToken()));
      assertFalse(filter.validateSignature(getInvalidToken()));
  
      verifyAll();
    }
  
    @Test
    public void testValidateAudiences() throws Exception {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
  
      assertTrue(filter.validateAudiences(getSignedToken()));
      assertFalse(filter.validateAudiences(getInvalidToken()));
  
      verifyAll();
    }
  
    @Test
    public void testValidateNullAudiences() throws Exception {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties(null)).anyTimes();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertTrue(filter.validateAudiences(getSignedToken()));
      assertTrue(filter.validateAudiences(getInvalidToken()));
  
      verifyAll();
    }
  
    @Test
    public void testValidateTokenWithoutAudiences() throws Exception {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertFalse(filter.validateAudiences(getSignedToken(null)));
  
      verifyAll();
    }
  
    @Test
    public void testValidateExpiration() throws Exception {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertTrue(filter.validateExpiration(getSignedToken()));
      assertFalse(filter.validateExpiration(getInvalidToken()));
  
      verifyAll();
    }
  
    @Test
    public void testValidateNoExpiration() throws Exception {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
  
      assertTrue(filter.validateExpiration(getSignedToken(null, "test-audience")));
      assertFalse(filter.validateExpiration(getInvalidToken()));
  
      verifyAll();
    }
  
    @Test
    public void testShouldApplyTrue() throws JOSEException {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      SignedJWT token = getInvalidToken();
  
      Cookie cookie = createMock(Cookie.class);
      expect(cookie.getName()).andReturn("non-default").atLeastOnce();
      expect(cookie.getValue()).andReturn(token.serialize()).atLeastOnce();
  
      HttpServletRequest request = createMock(HttpServletRequest.class);
      expect(request.getCookies()).andReturn(new Cookie[]{cookie});
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertTrue(filter.shouldApply(request));
  
      verifyAll();
    }
  
    @Test
    public void testShouldApplyTrueBadToken() throws JOSEException {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      Cookie cookie = createMock(Cookie.class);
      expect(cookie.getName()).andReturn("non-default").atLeastOnce();
      expect(cookie.getValue()).andReturn("bad token").atLeastOnce();
  
      HttpServletRequest request = createMock(HttpServletRequest.class);
      expect(request.getCookies()).andReturn(new Cookie[]{cookie});
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertTrue(filter.shouldApply(request));
  
      verifyAll();
    }
  
    @Test
    public void testShouldApplyFalseMissingCookie() throws JOSEException {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      Cookie cookie = createMock(Cookie.class);
      expect(cookie.getName()).andReturn("some-other-cookie").atLeastOnce();
  
      HttpServletRequest request = createMock(HttpServletRequest.class);
      expect(request.getCookies()).andReturn(new Cookie[]{cookie});
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertFalse(filter.shouldApply(request));
  
      verifyAll();
    }
  
    @Test
    public void testShouldApplyFalseNotEnabled() throws JOSEException {
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(null).anyTimes();
  
      HttpServletRequest request = createMock(HttpServletRequest.class);
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
  
      replayAll();
  
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(null, jwtAuthenticationPropertiesProvider, null, eventHandler);
      assertFalse(filter.shouldApply(request));
  
      verify(request);
    }
  
    @Test(expected = IllegalArgumentException.class)
    public void ensureNonNullEventHandler() {
      new AmbariJwtAuthenticationFilter(createNiceMock(AmbariEntryPoint.class), createNiceMock(JwtAuthenticationPropertiesProvider.class), createNiceMock(AmbariJwtAuthenticationProvider.class), null);
    }
  
    @Test
    public void testDoFilterSuccessful() throws Exception {
      Capture<? extends AmbariAuthenticationFilter> captureFilter = newCapture(CaptureType.ALL);
  
      SignedJWT token = getSignedToken();
  
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      Configuration configuration = createNiceMock(Configuration.class);
      expect(configuration.getMaxAuthenticationFailures()).andReturn(10).anyTimes();
  
      HttpServletRequest request = createMock(HttpServletRequest.class);
      HttpServletResponse response = createMock(HttpServletResponse.class);
      FilterChain filterChain = createMock(FilterChain.class);
  
      Cookie cookie = createMock(Cookie.class);
      expect(cookie.getName()).andReturn("non-default").once();
      expect(cookie.getValue()).andReturn(token.serialize()).once();
  
      expect(request.getCookies()).andReturn(new Cookie[]{cookie}).once();
  
      UserAuthenticationEntity userAuthenticationEntity = createMock(UserAuthenticationEntity.class);
      expect(userAuthenticationEntity.getAuthenticationType()).andReturn(UserAuthenticationType.JWT).anyTimes();
      expect(userAuthenticationEntity.getAuthenticationKey()).andReturn("").anyTimes();
  
      UserEntity userEntity = createMock(UserEntity.class);
      expect(userEntity.getAuthenticationEntities()).andReturn(Collections.singletonList(userAuthenticationEntity)).atLeastOnce();
  
      User user = createMock(User.class);
  
      Users users = createMock(Users.class);
      expect(users.getUserEntity("test-user")).andReturn(userEntity).once();
      expect(users.getUser(userEntity)).andReturn(user).once();
      expect(user.getUserName()).andReturn("test-user").atLeastOnce();
      expect(users.getUserAuthorities(userEntity)).andReturn(Collections.emptyList()).once();
      users.validateLogin(userEntity, "test-user");
      expectLastCall().once();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
      eventHandler.beforeAttemptAuthentication(capture(captureFilter), eq(request), eq(response));
      expectLastCall().once();
      eventHandler.onSuccessfulAuthentication(capture(captureFilter), eq(request), eq(response), anyObject(Authentication.class));
      expectLastCall().once();
  
      filterChain.doFilter(request, response);
      expectLastCall().once();
  
      AuthenticationEntryPoint entryPoint = createNiceMock(AmbariEntryPoint.class);
  
      replayAll();
  
      AmbariJwtAuthenticationProvider provider = new AmbariJwtAuthenticationProvider(users, configuration);
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(entryPoint, jwtAuthenticationPropertiesProvider, provider, eventHandler);
      filter.doFilter(request, response, filterChain);
  
      verifyAll();
  
      List<? extends AmbariAuthenticationFilter> capturedFilters = captureFilter.getValues();
      for (AmbariAuthenticationFilter capturedFiltered : capturedFilters) {
        assertSame(filter, capturedFiltered);
      }
    }
  
  
    @Test
    public void testDoFilterUnsuccessful() throws Exception {
      Capture<? extends AmbariAuthenticationFilter> captureFilter = newCapture(CaptureType.ALL);
  
      SignedJWT token = getSignedToken();
  
      Configuration configuration = createMock(Configuration.class);
  
      JwtAuthenticationPropertiesProvider jwtAuthenticationPropertiesProvider = createMock(JwtAuthenticationPropertiesProvider.class);
      expect(jwtAuthenticationPropertiesProvider.get()).andReturn(createTestProperties()).anyTimes();
  
      HttpServletRequest request = createMock(HttpServletRequest.class);
      HttpServletResponse response = createMock(HttpServletResponse.class);
  
      FilterChain filterChain = createMock(FilterChain.class);
  
      Cookie cookie = createMock(Cookie.class);
      expect(cookie.getName()).andReturn("non-default").once();
      expect(cookie.getValue()).andReturn(token.serialize()).once();
  
      expect(request.getCookies()).andReturn(new Cookie[]{cookie}).once();
  
      Users users = createMock(Users.class);
      expect(users.getUserEntity("test-user")).andReturn(null).once();
  
      AmbariAuthenticationEventHandler eventHandler = createNiceMock(AmbariAuthenticationEventHandler.class);
      eventHandler.beforeAttemptAuthentication(capture(captureFilter), eq(request), eq(response));
      expectLastCall().once();
      eventHandler.onUnsuccessfulAuthentication(capture(captureFilter), eq(request), eq(response), anyObject(AmbariAuthenticationException.class));
      expectLastCall().once();
  
      AuthenticationEntryPoint entryPoint = createNiceMock(AmbariEntryPoint.class);
      entryPoint.commence(eq(request), eq(response), anyObject(AmbariAuthenticationException.class));
      expectLastCall().once();
  
      replayAll();
  
      AmbariJwtAuthenticationProvider provider = new AmbariJwtAuthenticationProvider(users, configuration);
      AmbariJwtAuthenticationFilter filter = new AmbariJwtAuthenticationFilter(entryPoint, jwtAuthenticationPropertiesProvider, provider, eventHandler);
      filter.doFilter(request, response, filterChain);
  
      verifyAll();
  
      List<? extends AmbariAuthenticationFilter> capturedFilters = captureFilter.getValues();
      for (AmbariAuthenticationFilter capturedFiltered : capturedFilters) {
        assertSame(filter, capturedFiltered);
      }
    }
  
  }