import contextlib
import os
import subprocess
from subprocess import DEVNULL
import tempfile


# map principals to files storing credentials
_ccaches = {}


@contextlib.contextmanager
def krb5ccname_ctx(principal: str):
    """
    Temporarily set KRB5CCNAME to a ccache storing credentials
    for the specified user.
    """
    old_krb5ccname = os.environ['KRB5CCNAME']
    try:
        if principal not in _ccaches:
            f = tempfile.NamedTemporaryFile()
            os.environ['KRB5CCNAME'] = 'FILE:' + f.name
            args = ['kinit', principal]
            if principal == 'ceod/admin':
                args = ['kinit', '-k', principal]
            subprocess.run(
                args, stdout=DEVNULL, text=True, input='krb5',
                check=True)
            _ccaches[principal] = f
        else:
            os.environ['KRB5CCNAME'] = 'FILE:' + _ccaches[principal].name
        yield
    finally:
        os.environ['KRB5CCNAME'] = old_krb5ccname