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.cassandra.auth;
  
  import java.net.InetAddress;
  import java.net.UnknownHostException;
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.stream.Collectors;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.Lists;
  
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.cql3.CIDR;
  import org.apache.cassandra.cql3.QueryOptions;
  import org.apache.cassandra.cql3.QueryProcessor;
  import org.apache.cassandra.cql3.UntypedResultSet;
  import org.apache.cassandra.cql3.statements.SelectStatement;
  import org.apache.cassandra.db.ConsistencyLevel;
  import org.apache.cassandra.db.marshal.InetAddressType;
  import org.apache.cassandra.db.marshal.ShortType;
  import org.apache.cassandra.db.marshal.TupleType;
  import org.apache.cassandra.exceptions.RequestExecutionException;
  import org.apache.cassandra.schema.SchemaConstants;
  import org.apache.cassandra.service.ClientState;
  import org.apache.cassandra.transport.Dispatcher;
  import org.apache.cassandra.transport.messages.ResultMessage;
  import org.apache.cassandra.utils.ByteBufferUtil;
  import org.apache.cassandra.utils.MBeanWrapper;
  import org.apache.cassandra.utils.Pair;
  
  import static org.apache.cassandra.service.QueryState.forInternalCalls;
  
  public class CIDRGroupsMappingManager implements CIDRGroupsMappingManagerMBean
  {
      public static final String MBEAN_NAME = "org.apache.cassandra.db:type=CIDRGroupsMappingManager";
  
      private SelectStatement getCidrGroupsStatement;
      private SelectStatement getCidrsForCidrGroupStatement;
  
      private static final int PAGE_SIZE = 128;
  
      public void setup()
      {
          if (!MBeanWrapper.instance.isRegistered(MBEAN_NAME))
              MBeanWrapper.instance.registerMBean(this, MBEAN_NAME);
  
          String getCidrGroupsQuery = String.format("SELECT cidr_group FROM %s.%s",
                                                    SchemaConstants.AUTH_KEYSPACE_NAME,
                                                    AuthKeyspace.CIDR_GROUPS);
          getCidrGroupsStatement = (SelectStatement) QueryProcessor.getStatement(getCidrGroupsQuery,
                                                                                 ClientState.forInternalCalls());
  
          String getCidrsForCidrGroupQuery = String.format("SELECT cidrs FROM %s.%s where cidr_group = ?",
                                                           SchemaConstants.AUTH_KEYSPACE_NAME,
                                                           AuthKeyspace.CIDR_GROUPS);
          getCidrsForCidrGroupStatement = (SelectStatement) QueryProcessor.getStatement(getCidrsForCidrGroupQuery,
                                                                                        ClientState.forInternalCalls());
      }
  
      @VisibleForTesting
      ResultMessage.Rows select(SelectStatement statement, QueryOptions options)
      {
          return statement.execute(forInternalCalls(), options, Dispatcher.RequestTime.forImmediateExecution());
      }
  
      UntypedResultSet process(String query, ConsistencyLevel cl) throws RequestExecutionException
      {
          return QueryProcessor.process(query, cl);
      }
  
      protected static String getCidrTuplesSetString(List<CIDR> cidrs)
      {
          String inner = cidrs.stream().map(CIDR::asCqlTupleString)
                              .collect(Collectors.joining(", "));
          return '{' + inner + '}';
      }
  
      public UntypedResultSet getCidrGroupsTableEntries()
      {
          String getAllRowsQuery = String.format("SELECT * FROM %s.%s",
                                                 SchemaConstants.AUTH_KEYSPACE_NAME,
                                                 AuthKeyspace.CIDR_GROUPS);
  
          return QueryProcessor.executeInternalWithPaging(getAllRowsQuery, PAGE_SIZE);
      }
  
      public Set<String> getAvailableCidrGroups()
      {
          Set<String> availableCidrGroups = new HashSet<>();
  
          QueryOptions options = QueryOptions.forInternalCalls(CassandraAuthorizer.authReadConsistencyLevel(),
                                                               Collections.emptyList());
  
          ResultMessage.Rows rows = select(getCidrGroupsStatement, options);
          UntypedResultSet result = UntypedResultSet.create(rows.result);
  
          for (UntypedResultSet.Row row : result)
          {
              if (!row.has("cidr_group"))
                  throw new IllegalStateException("Invalid row " + row + " in table: " +
                                                  SchemaConstants.AUTH_KEYSPACE_NAME + '.' + AuthKeyspace.CIDR_GROUPS);
  
              availableCidrGroups.add(row.getString("cidr_group"));
          }
  
          return availableCidrGroups;
      }
  
      public Set<Pair<InetAddress, Short>> retrieveCidrsFromRow(UntypedResultSet.Row row)
      {
          if (!row.has("cidrs"))
              throw new RuntimeException("Invalid row, doesn't have column cidrs");
  
          Set<Pair<InetAddress, Short>> cidrs = new HashSet<>();
          TupleType tupleType = new TupleType(Arrays.asList(InetAddressType.instance, ShortType.instance));
  
          Set<ByteBuffer> cidrAsTuples = row.getFrozenSet("cidrs", tupleType);
          for (ByteBuffer cidrAsTuple : cidrAsTuples)
          {
              List<ByteBuffer> elements = tupleType.unpack(cidrAsTuple);
              InetAddress ip = InetAddressType.instance.compose(elements.get(0));
              short netMask = ShortType.instance.compose(elements.get(1));
              cidrs.add(Pair.create(ip, netMask));
          }
  
          return cidrs;
      }
  
      public Set<String> getCidrsOfCidrGroupAsStrings(String cidrGroupName)
      {
          QueryOptions options = QueryOptions.forInternalCalls(CassandraAuthorizer.authReadConsistencyLevel(),
                                                               Lists.newArrayList(ByteBufferUtil.bytes(cidrGroupName)));
          ResultMessage.Rows rows = select(getCidrsForCidrGroupStatement, options);
          UntypedResultSet result = UntypedResultSet.create(rows.result);
  
          if (result.isEmpty())
              return Collections.emptySet();
  
          Set<Pair<InetAddress, Short>> allCidrs = retrieveCidrsFromRow(result.one());
          Set<String> cidrStrs = new HashSet<>();
          for (Pair<InetAddress, Short> cidr : allCidrs)
              cidrStrs.add(cidr.left().getHostAddress() + '/' + cidr.right());
  
          return cidrStrs;
      }
  
      public void updateCidrGroup(String cidrGroupName, List<String> cidrs)
      {
          List<CIDR> validCidrs = new ArrayList<>();
          for (String cidr : cidrs)
          {
              validCidrs.add(CIDR.getInstance(cidr));
          }
  
          String query = String.format("UPDATE %s.%s SET cidrs = %s WHERE cidr_group = '%s'",
                                       SchemaConstants.AUTH_KEYSPACE_NAME,
                                       AuthKeyspace.CIDR_GROUPS,
                                       getCidrTuplesSetString(validCidrs),
                                       cidrGroupName);
  
          process(query, CassandraAuthorizer.authWriteConsistencyLevel());
      }
  
      @VisibleForTesting
      void dropCidrGroupIfExists(String cidrGroupName)
      {
          String query = String.format("DELETE FROM %s.%s WHERE cidr_group = '%s'",
                                       SchemaConstants.AUTH_KEYSPACE_NAME,
                                       AuthKeyspace.CIDR_GROUPS,
                                       cidrGroupName);
  
          process(query, CassandraAuthorizer.authWriteConsistencyLevel());
      }
  
      public void dropCidrGroup(String cidrGroupName)
      {
          Set<String> cidrs = getCidrsOfCidrGroupAsStrings(cidrGroupName);
          if (cidrs.isEmpty())
              throw new RuntimeException("CIDR group '" + cidrGroupName + "' doesn't exists");
  
          dropCidrGroupIfExists(cidrGroupName);
      }
  
      public void recreateCidrGroupsMapping(Map<String, List<String>> cidrGroupsMapping)
      {
          Set<String> existingMappings = getAvailableCidrGroups();
  
          // Overwrites mappings of existing cidr groups and inserts new cidr groups
          for (Map.Entry<String, List<String>> cidrGroupMapping : cidrGroupsMapping.entrySet())
          {
              String cidrGroupName = cidrGroupMapping.getKey();
              updateCidrGroup(cidrGroupName, cidrGroupMapping.getValue());
              existingMappings.remove(cidrGroupName);
          }
  
          // Delete old CIDR groups which do not exist in new mappings
          for (String cidrGroupName : existingMappings)
          {
              dropCidrGroupIfExists(cidrGroupName);
          }
      }
  
      public Set<String> getCidrGroupsOfIP(String ipStr)
      {
          try
          {
              return DatabaseDescriptor.getCIDRAuthorizer().lookupCidrGroupsForIp(InetAddress.getByName(ipStr));
          }
          catch (UnknownHostException e)
          {
              throw new IllegalArgumentException("Invalid IP " + ipStr, e);
          }
      }
  
      public void loadCidrGroupsCache()
      {
          DatabaseDescriptor.getCIDRAuthorizer().loadCidrGroupsCache();
      }
  }
  