# Implementation of GSMA eSIM RSP (Remote SIM Provisioning) ES8+
# as per SGP22 v3.0 Section 5.5
#
# (C) 2023-2024 by Harald Welte <laforge@osmocom.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from typing import Dict, List, Optional
from pySim.utils import b2h, h2b, bertlv_encode_tag, bertlv_encode_len

import pySim.esim.rsp as rsp
from pySim.esim.bsp import BspInstance

# Given that GSMA RSP uses ASN.1 in a very weird way, we actually cannot encode the full data type before
# signing, but we have to build parts of it separately first, then sign that, so we can put the signature
# into the same sequence as the signed data.  We use the existing pySim TLV code for this.

def wrap_as_der_tlv(tag: int, val: bytes) -> bytes:
    """Wrap the 'value' into a DER-encoded TLV."""
    return bertlv_encode_tag(tag) + bertlv_encode_len(len(val)) + val

def gen_init_sec_chan_signed_part(iscsp: Dict) -> bytes:
    """Generate the concatenated remoteOpId, transactionId, controlRefTemplate and smdpOtpk data objects
    without the outer SEQUENCE tag / length or the remainder of initialiseSecureChannel, as is required
    for signing purpose."""
    out = b''
    out += wrap_as_der_tlv(0x82, bytes([iscsp['remoteOpId']]))
    out += wrap_as_der_tlv(0x80, iscsp['transactionId'])

    crt = iscsp['controlRefTemplate']
    out_crt = wrap_as_der_tlv(0x80, crt['keyType'])
    out_crt += wrap_as_der_tlv(0x81, crt['keyLen'])
    out_crt += wrap_as_der_tlv(0x84, crt['hostId'])
    out += wrap_as_der_tlv(0xA6, out_crt)

    out += wrap_as_der_tlv(0x5F49, iscsp['smdpOtpk'])
    return out


# SGP.22 Section 5.5.1
def gen_initialiseSecureChannel(transactionId: str, host_id: bytes, smdp_otpk: bytes, euicc_otpk: bytes, dp_pb):
    """Generate decoded representation of (signed) initialiseSecureChannel (SGP.22 5.5.2)"""
    init_scr = { 'remoteOpId': 1, # installBoundProfilePackage
                 'transactionId': h2b(transactionId),
                 # GlobalPlatform Card Specification Amendment F [13] section 6.5.2.3 for the Mutual Authentication Data Field
                 'controlRefTemplate': { 'keyType': bytes([0x88]), 'keyLen': bytes([16]), 'hostId': host_id },
                 'smdpOtpk': smdp_otpk, # otPK.DP.KA
                }
    to_sign = gen_init_sec_chan_signed_part(init_scr) + wrap_as_der_tlv(0x5f49, euicc_otpk)
    init_scr['smdpSign'] = dp_pb.ecdsa_sign(to_sign)
    return init_scr

def gen_replace_session_keys(ppk_enc: bytes, ppk_cmac: bytes, initial_mcv: bytes) -> bytes:
    """Generate encoded (but unsigned) ReplaceSessionKeysReqest DO (SGP.22 5.5.4)"""
    rsk = { 'ppkEnc': ppk_enc, 'ppkCmac': ppk_cmac, 'initialMacChainingValue': initial_mcv }
    return rsp.asn1.encode('ReplaceSessionKeysRequest', rsk)


class ProfileMetadata:
    """Representation of Profile metadata. Right now only the mandatory bits are
    supported, but in general this should follow the StoreMetadataRequest of SGP.22 5.5.3"""
    def __init__(self, iccid_bin: bytes, spn: str, profile_name: str):
        self.iccid_bin = iccid_bin
        self.spn = spn
        self.profile_name = profile_name

    def gen_store_metadata_request(self) -> bytes:
        """Generate encoded (but unsigned) StoreMetadataReqest DO (SGP.22 5.5.3)"""
        smr = {
            'iccid': self.iccid_bin,
            'serviceProviderName': self.spn,
            'profileName': self.profile_name,
        }
        return rsp.asn1.encode('StoreMetadataRequest', smr)


class ProfilePackage:
    def __init__(self, metadata: Optional[ProfileMetadata] = None):
        self.metadata = metadata

class UnprotectedProfilePackage(ProfilePackage):
    """Representing an unprotected profile package (UPP) as defined in SGP.22 Section 2.5.2"""

    @classmethod
    def from_der(cls, der: bytes, metadata: Optional[ProfileMetadata] = None) -> 'UnprotectedProfilePackage':
        """Load an UPP from its DER representation."""
        inst = cls(metadata=metadata)
        cls.der = der
        # TODO: we later certainly want to parse it so we can perform modification (IMSI, key material, ...)
        # just like in the traditional SIM/USIM dynamic data phase at the end of personalization
        return inst

    def to_der(self):
        """Return the DER representation of the UPP."""
        # TODO: once we work on decoded structures, we may want to re-encode here
        return self.der

class ProtectedProfilePackage(ProfilePackage):
    """Representing a protected profile package (PPP) as defined in SGP.22 Section 2.5.3"""

    @classmethod
    def from_upp(cls, upp: UnprotectedProfilePackage, bsp: BspInstance) -> 'ProtectedProfilePackage':
        """Generate the PPP as a sequence of encrypted and MACed Command TLVs representing the UPP"""
        inst = cls(metadata=upp.metadata)
        inst.upp = upp
        # store ppk-enc, ppc-mac
        inst.ppk_enc = bsp.c_algo.s_enc
        inst.ppk_mac = bsp.m_algo.s_mac
        inst.initial_mcv = bsp.m_algo.mac_chain
        inst.encoded = bsp.encrypt_and_mac(0x86, upp.to_der())
        return inst

    #def __val__(self):
        #return self.encoded

class BoundProfilePackage(ProfilePackage):
    """Representing a bound profile package (BPP) as defined in SGP.22 Section 2.5.4"""

    @classmethod
    def from_ppp(cls, ppp: ProtectedProfilePackage):
        inst = cls()
        inst.upp = None
        inst.ppp = ppp
        return inst

    @classmethod
    def from_upp(cls, upp: UnprotectedProfilePackage):
        inst = cls()
        inst.upp = upp
        inst.ppp = None
        return inst

    def encode(self, ss: 'RspSessionState', dp_pb: 'CertAndPrivkey') -> bytes:
        """Generate a bound profile package (SGP.22 2.5.4)."""

        def encode_seq(tag: int, sequence: List[bytes]) -> bytes:
            """Encode a "sequenceOfXX" as specified in SGP.22 specifying the raw SEQUENCE OF tag,
            and assuming the caller provides the fully-encoded (with TAG + LEN) member TLVs."""
            payload = b''.join(sequence)
            return bertlv_encode_tag(tag) + bertlv_encode_len(len(payload)) + payload

        bsp = BspInstance.from_kdf(ss.shared_secret, 0x88, 16, ss.host_id, h2b(ss.eid))

        iscr = gen_initialiseSecureChannel(ss.transactionId, ss.host_id, ss.smdp_otpk, ss.euicc_otpk, dp_pb)
        # generate unprotected input data
        conf_idsp_bin = rsp.asn1.encode('ConfigureISDPRequest', {})
        if self.upp:
            smr_bin = self.upp.metadata.gen_store_metadata_request()
        else:
            smr_bin = self.ppp.metadata.gen_store_metadata_request()

        # we don't use rsp.asn1.encode('boundProfilePackage') here, as the BSP already provides
        # fully encoded + MACed TLVs including their tag + length values.  We cannot put those as
        # 'value' input into an ASN.1 encoder, as that would double the TAG + LENGTH :(

        # 'initialiseSecureChannelRequest'
        bpp_seq = rsp.asn1.encode('InitialiseSecureChannelRequest', iscr)
        # firstSequenceOf87
        bpp_seq += encode_seq(0xa0, bsp.encrypt_and_mac(0x87, conf_idsp_bin))
        # sequenceOF88
        bpp_seq += encode_seq(0xa1, bsp.mac_only(0x88, smr_bin))

        if self.ppp: # we have to use session keys
            rsk_bin = gen_replace_session_keys(self.ppp.ppk_enc, self.ppp.ppk_mac, self.ppp.initial_mcv)
            # secondSequenceOf87
            bpp_seq += encode_seq(0xa2, bsp.encrypt_and_mac(0x87, rsk_bin))
        else:
            self.ppp = ProtectedProfilePackage.from_upp(self.upp, bsp)

        # 'sequenceOf86'
        bpp_seq += encode_seq(0xa3, self.ppp.encoded)

        # manual DER encode: wrap in outer SEQUENCE
        return bertlv_encode_tag(0xbf36) + bertlv_encode_len(len(bpp_seq)) + bpp_seq