/*
 * Copyright (C) 2022 The Android Open Source Project
 *
 * Licensed 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 com.android.server.wifi;

import android.net.MacAddress;
import android.net.wifi.WifiConfiguration;
import android.os.Handler;
import android.util.Log;
import android.util.SparseArray;

import com.android.internal.annotations.VisibleForTesting;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/** Utilities for storing PMK cache. */
public class PmkCacheManager {
    private static final String TAG = "PmkCacheManager";

    @VisibleForTesting
    static final String PMK_CACHE_EXPIRATION_ALARM_TAG = "PMK_CACHE_EXPIRATION_TIMER";

    private final Clock mClock;
    private final Handler mEventHandler;

    private boolean mVerboseLoggingEnabled = false;
    private SparseArray<List<PmkCacheStoreData>> mPmkCacheEntries = new SparseArray<>();

    public PmkCacheManager(Clock clock, Handler eventHandler) {
        mClock = clock;
        mEventHandler = eventHandler;
    }

    /**
     * Add a PMK cache entry to the store.
     *
     * @param macAddress the interface MAC address to connect to the network.
     * @param networkId the network ID of the WifiConfiguration associates with the network.
     * @param expirationTimeInSec the expiration time of the PMK cache since boot.
     * @param serializedEntry the opaque data of the PMK cache.
     * @return true when PMK cache is added; otherwise, false.
     */
    public boolean add(MacAddress macAddress, int networkId,
            long expirationTimeInSec, ArrayList<Byte> serializedEntry) {
        if (WifiConfiguration.INVALID_NETWORK_ID == networkId) return false;
        if (macAddress == null) {
            Log.w(TAG, "Omit PMK cache due to no valid MAC address");
            return false;
        }
        if (null == serializedEntry) {
            Log.w(TAG, "Omit PMK cache due to null entry.");
            return false;
        }
        final long elapseTimeInSecond = mClock.getElapsedSinceBootMillis() / 1000;
        if (elapseTimeInSecond >= expirationTimeInSec) {
            Log.w(TAG, "Omit expired PMK cache.");
            return false;
        }

        PmkCacheStoreData newStoreData =
                new PmkCacheStoreData(macAddress, serializedEntry, expirationTimeInSec);
        List<PmkCacheStoreData> pmkDataList = mPmkCacheEntries.get(networkId);
        if (pmkDataList == null) {
            pmkDataList = new ArrayList<>();
            mPmkCacheEntries.put(networkId, pmkDataList);
        } else {
            PmkCacheStoreData existStoreData = pmkDataList.stream()
                    .filter(storeData -> Objects.equals(storeData, newStoreData))
                    .findAny()
                    .orElse(null);
            if (null != existStoreData) {
                if (mVerboseLoggingEnabled) {
                    Log.d(TAG, "PMK entry exists, skip it.");
                }
                return true;
            }
        }

        pmkDataList.add(newStoreData);
        if (mVerboseLoggingEnabled) {
            Log.d(TAG, "Network " + networkId + " PmkCache Count: " + pmkDataList.size());
        }
        updatePmkCacheExpiration();
        return true;
    }

    /**
     * Remove PMK caches associated with the network ID.
     *
     * @param networkId the network ID of PMK caches to be removed.
     * @return true when PMK caches are removed; otherwise, false.
     */
    public boolean remove(int networkId) {
        if (WifiConfiguration.INVALID_NETWORK_ID == networkId) return false;
        if (!mPmkCacheEntries.contains(networkId)) return false;

        mPmkCacheEntries.remove(networkId);
        updatePmkCacheExpiration();
        return true;
    }

    /**
     * Remove PMK caches associated with the network ID when the interface
     * MAC address is changed.
     *
     * @param networkId the network ID of PMK caches to be removed.
     * @param curMacAddress current interface MAC address.
     * @return true when PMK caches are removed; otherwise, false.
     */

    public boolean remove(int networkId, MacAddress curMacAddress) {
        if (WifiConfiguration.INVALID_NETWORK_ID == networkId) return false;
        List<PmkCacheStoreData> pmkDataList = mPmkCacheEntries.get(networkId);
        if (null == pmkDataList) return false;

        pmkDataList.removeIf(pmkData -> !Objects.equals(curMacAddress, pmkData.macAddress));

        if (pmkDataList.size() == 0) {
            remove(networkId);
        }
        return true;
    }

    /**
     * Get PMK caches associated with the network ID.
     *
     * @param networkId the network ID to be queried.
     * @return A list of PMK caches associated with the network ID.
     *         If none of PMK cache is associated with the network ID, return null.
     */
    public List<ArrayList<Byte>> get(int networkId) {
        List<PmkCacheStoreData> pmkDataList = mPmkCacheEntries.get(networkId);
        if (WifiConfiguration.INVALID_NETWORK_ID == networkId) return null;
        if (null == pmkDataList) return null;

        final long elapseTimeInSecond = mClock.getElapsedSinceBootMillis() / 1000;
        List<ArrayList<Byte>> dataList = new ArrayList<>();
        for (PmkCacheStoreData pmkData: pmkDataList) {
            if (pmkData.isValid(elapseTimeInSecond)) {
                dataList.add(pmkData.data);
            }
        }
        return dataList;
    }

    /**
     * Enable/Disable verbose logging.
     *
     * @param verboseEnabled Verbose flag set in overlay XML.
     */
    public void enableVerboseLogging(boolean verboseEnabled) {
        mVerboseLoggingEnabled = verboseEnabled;
    }

    @VisibleForTesting
    void updatePmkCacheExpiration() {
        mEventHandler.removeCallbacksAndMessages(PMK_CACHE_EXPIRATION_ALARM_TAG);

        long elapseTimeInSecond = mClock.getElapsedSinceBootMillis() / 1000;
        long nextUpdateTimeInSecond = Long.MAX_VALUE;
        if (mVerboseLoggingEnabled) {
            Log.d(TAG, "Update PMK cache expiration at " + elapseTimeInSecond);
        }

        List<Integer> emptyStoreDataList = new ArrayList<>();
        for (int i = 0; i < mPmkCacheEntries.size(); i++) {
            int networkId = mPmkCacheEntries.keyAt(i);
            List<PmkCacheStoreData> list = mPmkCacheEntries.get(networkId);
            list.removeIf(pmkData -> !pmkData.isValid(elapseTimeInSecond));
            if (list.size() == 0) {
                emptyStoreDataList.add(networkId);
                continue;
            }
            for (PmkCacheStoreData pmkData: list) {
                if (nextUpdateTimeInSecond > pmkData.expirationTimeInSec) {
                    nextUpdateTimeInSecond = pmkData.expirationTimeInSec;
                }
            }
        }
        emptyStoreDataList.forEach(networkId -> mPmkCacheEntries.remove(networkId));

        // No need to arrange next update since there is no valid PMK in the cache.
        if (nextUpdateTimeInSecond == Long.MAX_VALUE) {
            return;
        }

        if (mVerboseLoggingEnabled) {
            Log.d(TAG, "PMK cache next expiration time: " + nextUpdateTimeInSecond);
        }
        long delayedTimeInMs = (nextUpdateTimeInSecond - elapseTimeInSecond) * 1000;
        mEventHandler.postDelayed(
                () -> {
                    updatePmkCacheExpiration();
                },
                PMK_CACHE_EXPIRATION_ALARM_TAG,
                (delayedTimeInMs > 0) ? delayedTimeInMs : 0);
    }

    private static class PmkCacheStoreData {

        public MacAddress macAddress;
        public ArrayList<Byte> data;
        public long expirationTimeInSec;

        PmkCacheStoreData(MacAddress macAddr, ArrayList<Byte> serializedData, long timeInSec) {
            macAddress = macAddr;
            data = serializedData;
            expirationTimeInSec = timeInSec;
        }

        /**
         * Validate this PMK cache against the timestamp.
         *
         * @param currentTimeInSec the timestamp to be checked.
         * @return true if this PMK cache is valid against the timestamp; otherwise, false.
         */
        public boolean isValid(long currentTimeInSec) {
            return expirationTimeInSec > 0 && expirationTimeInSec > currentTimeInSec;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (!(o instanceof PmkCacheStoreData)) return false;
            PmkCacheStoreData storeData = (PmkCacheStoreData) o;
            return expirationTimeInSec == storeData.expirationTimeInSec
                    && Objects.equals(macAddress, storeData.macAddress)
                    && Objects.equals(data, storeData.data);
        }

        @Override
        public int hashCode() {
            return Objects.hash(macAddress, data, expirationTimeInSec);
        }
    }
}