import ctypes
from ctypes import wintypes

from .api import security
from .api import privilege
from .api import process


def get_process_token():
    """
    Get the current process token
    """
    token = wintypes.HANDLE()
    res = process.OpenProcessToken(
        process.GetCurrentProcess(), process.TOKEN_ALL_ACCESS, token
    )
    if not res > 0:
        raise RuntimeError("Couldn't get process token")
    return token


def get_symlink_luid():
    """
    Get the LUID for the SeCreateSymbolicLinkPrivilege
    """
    symlink_luid = privilege.LUID()
    res = privilege.LookupPrivilegeValue(
        None, "SeCreateSymbolicLinkPrivilege", symlink_luid
    )
    if not res > 0:
        raise RuntimeError("Couldn't lookup privilege value")
    return symlink_luid


def get_privilege_information():
    """
    Get all privileges associated with the current process.
    """
    # first call with zero length to determine what size buffer we need

    return_length = wintypes.DWORD()
    params = [
        get_process_token(),
        privilege.TOKEN_INFORMATION_CLASS.TokenPrivileges,
        None,
        0,
        return_length,
    ]

    res = privilege.GetTokenInformation(*params)

    # assume we now have the necessary length in return_length

    buffer = ctypes.create_string_buffer(return_length.value)
    params[2] = buffer
    params[3] = return_length.value

    res = privilege.GetTokenInformation(*params)
    assert res > 0, "Error in second GetTokenInformation (%d)" % res

    privileges = ctypes.cast(
        buffer, ctypes.POINTER(privilege.TOKEN_PRIVILEGES)
    ).contents
    return privileges


def report_privilege_information():
    """
    Report all privilege information assigned to the current process.
    """
    privileges = get_privilege_information()
    print("found {0} privileges".format(privileges.count))
    tuple(map(print, privileges))


def enable_symlink_privilege():
    """
    Try to assign the symlink privilege to the current process token.
    Return True if the assignment is successful.
    """
    # create a space in memory for a TOKEN_PRIVILEGES structure
    #  with one element
    size = ctypes.sizeof(privilege.TOKEN_PRIVILEGES)
    size += ctypes.sizeof(privilege.LUID_AND_ATTRIBUTES)
    buffer = ctypes.create_string_buffer(size)
    tp = ctypes.cast(buffer, ctypes.POINTER(privilege.TOKEN_PRIVILEGES)).contents
    tp.count = 1
    tp.get_array()[0].enable()
    tp.get_array()[0].LUID = get_symlink_luid()
    token = get_process_token()
    res = privilege.AdjustTokenPrivileges(token, False, tp, 0, None, None)
    if res == 0:
        raise RuntimeError("Error in AdjustTokenPrivileges")

    ERROR_NOT_ALL_ASSIGNED = 1300
    return ctypes.windll.kernel32.GetLastError() != ERROR_NOT_ALL_ASSIGNED


class PolicyHandle(wintypes.HANDLE):
    pass


class LSA_UNICODE_STRING(ctypes.Structure):
    _fields_ = [
        ('length', ctypes.c_ushort),
        ('max_length', ctypes.c_ushort),
        ('buffer', ctypes.wintypes.LPWSTR),
    ]


def OpenPolicy(system_name, object_attributes, access_mask):
    policy = PolicyHandle()
    raise NotImplementedError(
        "Need to construct structures for parameters "
        "(see http://msdn.microsoft.com/en-us/library/windows"
        "/desktop/aa378299%28v=vs.85%29.aspx)"
    )
    res = ctypes.windll.advapi32.LsaOpenPolicy(
        system_name, object_attributes, access_mask, ctypes.byref(policy)
    )
    assert res == 0, "Error status {res}".format(**vars())
    return policy


def grant_symlink_privilege(who, machine=''):
    """
    Grant the 'create symlink' privilege to who.

    Based on http://support.microsoft.com/kb/132958
    """
    flags = security.POLICY_CREATE_ACCOUNT | security.POLICY_LOOKUP_NAMES
    policy = OpenPolicy(machine, flags)
    return policy


def main():
    assigned = enable_symlink_privilege()
    msg = ['failure', 'success'][assigned]

    print("Symlink privilege assignment completed with {0}".format(msg))


if __name__ == '__main__':
    main()