import sys
import threading
from typing import Dict, Optional, Set, Tuple

from keylime.common.algorithms import Hash
from keylime.ima.file_signatures import ImaKeyrings

if sys.version_info >= (3, 7):
    from dataclasses import dataclass
else:
    from keylime.backport_dataclasses import dataclass


@dataclass
class TPMClockInfo:
    clock: int
    resetcount: int
    restartcount: int
    safe: int

    @classmethod
    def from_dict(cls, dclki: Dict[str, int]) -> "TPMClockInfo":
        return cls(
            clock=dclki.get("clock", 0),
            resetcount=dclki.get("resetCount", 0),
            restartcount=dclki.get("restartCount", 0),
            safe=dclki.get("safe", 1),
        )

    def to_dict(self) -> Dict[str, int]:
        return {
            "clock": self.clock,
            "resetCount": self.resetcount,
            "restartCount": self.restartcount,
            "safe": self.safe,
        }


class TPMState:
    """TPMState models the state of the TPM's PCRs"""

    pcrs: Dict[int, Optional[bytes]]
    hash_alg: Dict[int, Hash]

    def __init__(self) -> None:
        """constructor"""
        self.pcrs = {}
        self.hash_alg = {}  # Record the hash algorithm that a given PCR uses
        for pcr_num in range(0, 24):
            self.reset_pcr(pcr_num)

    def init_pcr(self, pcr_num: int, hash_alg: Hash) -> None:
        """Initializes a PCR"""
        if pcr_num not in self.hash_alg:
            self.hash_alg[pcr_num] = hash_alg

        if self.pcrs[pcr_num] is None:
            if 17 <= pcr_num <= 23:
                self.pcrs[pcr_num] = self.hash_alg[pcr_num].get_ff_hash()
            else:
                self.pcrs[pcr_num] = self.hash_alg[pcr_num].get_start_hash()

    def reset_pcr(self, pcr_num: int) -> None:
        """Reset a specific PCR"""
        self.pcrs[pcr_num] = None

    def get_pcr(self, pcr_num: int) -> Optional[bytes]:
        """
        Get the state of a PCR.

        :returns: PCR value or None if not initialized
        """
        return self.pcrs.get(pcr_num, None)

    def used_pcr(self, pcr_num: int) -> bool:
        """Check if a PCR was actually requested at least once"""
        return (pcr_num in self.hash_alg) and (self.pcrs[pcr_num] is not None)

    def set_pcr(self, pcr_num: int, pcr_value: bytes) -> None:
        """Set the value of a PCR"""
        self.pcrs[pcr_num] = pcr_value


class AgentAttestState:
    """AgentAttestState is used to support incremental attestation"""

    agent_id: str
    next_ima_ml_entry: int
    boottime: int
    tpm_clockinfo: TPMClockInfo
    tpm_state: TPMState
    ima_pcrs: Set[int]
    ima_keyring: ImaKeyrings
    ima_dm_state: Optional[bytes]
    quote_progress: Optional[Tuple[int, int]]

    def __init__(self, agent_id: str) -> None:
        """constructor"""

        self.agent_id = agent_id
        self.next_ima_ml_entry = 0
        self.set_boottime(0)
        self.tpm_clockinfo = TPMClockInfo(clock=0, resetcount=0, restartcount=0, safe=1)

        self.tpm_state = TPMState()
        self.ima_pcrs = set()

        self.ima_keyrings = ImaKeyrings()

        self.reset_ima_attestation()

        self.ima_dm_state = None

        # quote_progress later consists of tuple(pcr_match_line, log length)
        self.quote_progress = None

    def get_agent_id(self) -> str:
        """Get the agent_id"""
        return self.agent_id

    def reset_ima_attestation(self) -> None:
        """Reset the IMA attestation state to start over with 1st entry
        and start over with learning the keys"""
        self.next_ima_ml_entry = 0
        for pcr_num in self.ima_pcrs:
            self.tpm_state.reset_pcr(pcr_num)
        self.set_boottime(0)
        self.ima_keyrings = ImaKeyrings()
        self.quote_progress = None

    def update_ima_attestation(self, pcr_num: int, pcr_value: bytes, num_ml_entries: int) -> None:
        """Update the attestation by remembering the new PCR value and the
        number of lines that were successfully processed."""
        self.ima_pcrs.add(pcr_num)
        self.tpm_state.set_pcr(pcr_num, pcr_value)
        self.next_ima_ml_entry += num_ml_entries

    def get_ima_pcrs(self) -> Dict[int, Optional[bytes]]:
        """Return a dict with the IMA pcrs"""
        ima_pcrs_dict = {}
        for pcr_num in self.ima_pcrs:
            # Only output IMA PCRs that were used at least once
            if self.tpm_state.used_pcr(pcr_num):
                ima_pcrs_dict[pcr_num] = self.tpm_state.get_pcr(pcr_num)
        return ima_pcrs_dict

    def set_ima_pcrs(self, ima_pcrs_dict: Dict[int, bytes]) -> None:
        """Set the values of the given ima_pcrs dict in the tpm_state"""
        for pcr_num, pcr_value in ima_pcrs_dict.items():
            self.tpm_state.set_pcr(pcr_num, pcr_value)
        self.ima_pcrs = set(ima_pcrs_dict.keys())

    def get_next_ima_ml_entry(self) -> int:
        """Return the next IMA measurement list entry we want to request from agent"""
        return self.next_ima_ml_entry

    def set_next_ima_ml_entry(self, next_ima_ml_entry: int) -> None:
        """Set the value of the next_ima_ml_entry field"""
        self.next_ima_ml_entry = next_ima_ml_entry

    def get_pcr_state(self, pcr_num: int, hash_alg: Hash = Hash.SHA1) -> Optional[bytes]:
        """Return the PCR state of the given PCR"""
        if not self.tpm_state.used_pcr(pcr_num):
            self.tpm_state.init_pcr(pcr_num, hash_alg)
        return self.tpm_state.get_pcr(pcr_num)

    def get_boottime(self) -> int:
        """Return the boottime of the system"""
        return self.boottime

    def set_boottime(self, boottime: int) -> None:
        """Set the boottime of the system"""
        self.boottime = boottime

    def get_tpm_clockinfo(self) -> TPMClockInfo:
        """Return the clock info extracted from a TPM quote"""
        return self.tpm_clockinfo

    def set_tpm_clockinfo(self, tpm_clockinfo: TPMClockInfo) -> None:
        """Set the clock info with information extracted from a TPM quote"""
        self.tpm_clockinfo = tpm_clockinfo

    def is_expected_boottime(self, boottime: int) -> bool:
        """Check whether the given boottime is the expected boottime"""
        return self.boottime == boottime

    def set_ima_keyrings(self, ima_keyrings: ImaKeyrings) -> None:
        """Set the ImaKeyrings object"""
        self.ima_keyrings = ima_keyrings

    def get_ima_keyrings(self) -> ImaKeyrings:
        """Get the ImaKeyrings object"""
        return self.ima_keyrings

    def get_ima_dm_state(self) -> Optional[bytes]:
        """Get encoded state of the DmValidator"""
        return self.ima_dm_state

    def set_ima_dm_state(self, state: bytes) -> None:
        self.ima_dm_state = state

    def check_quote_progress(self, pcr_match_line: int, log_length: int) -> bool:
        """Check whether the quote is making forward-progress.
        Return False if the previous quote has not caught up with the log
        from the previous report (previous lag).
        pcr_match_line must be a number >= 0 and <= log_length. An empty log
        has length 0. If the quote matches the state of the PCRs without
        replaying any further log entries, then this match occurs at
        pcr_match_line = 0 (which is at the same log entry as before). If 'n'
        lines were replayed, then the match occurred at line 'n'.
        """

        if not self.quote_progress:
            # set initial state
            self.quote_progress = (pcr_match_line, log_length)
            return True

        prev_pcr_match_line, prev_log_length = self.quote_progress
        self.quote_progress = (pcr_match_line, log_length)

        # quote caught up with log?
        if pcr_match_line == log_length:
            return True

        prev_lag = prev_log_length - prev_pcr_match_line

        # Insufficient quote forward-progress means that the quote was lagging
        # before and now the PCR(s) did not catch up with the previous lag.
        if prev_lag > 0 and pcr_match_line < prev_lag:
            return False
        return True


class AgentAttestStates:
    """AgentAttestStates administers a map of AgentAttestState's indexed by agent_id"""

    instance = None
    map_lock: threading.Lock
    map: Dict[str, AgentAttestState]

    @staticmethod
    def get_instance() -> "AgentAttestStates":
        """Create and return a singleton AgentAttestState"""
        if not AgentAttestStates.instance:
            AgentAttestStates.instance = AgentAttestStates()
        return AgentAttestStates.instance

    def __init__(self) -> None:
        """constructor"""
        self.map_lock = threading.Lock()
        self.map = {}

    def get_by_agent_id(self, agent_id: str) -> AgentAttestState:
        """Get an agent's state given its id"""

        with self.map_lock:
            agent_attest_state = self.map.get(agent_id)
            if not agent_attest_state:
                agent_attest_state = AgentAttestState(agent_id)
                self.map[agent_id] = agent_attest_state

        return agent_attest_state

    def delete_by_agent_id(self, agent_id: str) -> None:
        """Delete an agent's state given its id"""

        with self.map_lock:
            try:
                del self.map[agent_id]
            except KeyError:
                pass

    def add(
        self,
        agent_id: str,
        boottime: int,
        ima_pcrs_dict: Dict[int, bytes],
        next_ima_ml_entry: int,
        learned_ima_keyrings: Dict[str, str],
    ) -> None:
        """Add or replace an existing AgentAttestState initialized with the given values"""
        agentAttestState = self.get_by_agent_id(agent_id)
        agentAttestState.set_boottime(boottime)
        agentAttestState.set_ima_pcrs(ima_pcrs_dict)
        agentAttestState.set_next_ima_ml_entry(next_ima_ml_entry)
        agentAttestState.set_ima_keyrings(ImaKeyrings.from_json(learned_ima_keyrings))
