selfprivacy-rest-api/selfprivacy_api/graphql/common_types/user.py

import typing
import strawberry
from selfprivacy_api.graphql.mutations.mutation_interface import (
    MutationReturnInterface,
)
from enum import Enum

from selfprivacy_api.utils import ReadUserData


@strawberry.enum
class UserType(Enum):
    NORMAL = "NORMAL"
    PRIMARY = "PRIMARY"
    ROOT = "ROOT"


@strawberry.type
class User:
    """Users management"""

    user_type: UserType
    username: str
    # userHomeFolderspace: UserHomeFolderUsage
    ssh_keys: typing.List[str] = strawberry.field(default_factory=list)


@strawberry.type
class UserMutationReturn(MutationReturnInterface):
    """Return type for user mutation"""

    user: typing.Optional[User]

def get_user_by_username(username: str) -> typing.Optional[User]:
    with ReadUserData() as data:

        if username == "root":
            if data["ssh"]["rootKeys"] not in data:
                data["ssh"]["rootKeys"] = []

            return User(
                user_type=UserType.ROOT,
                username="root",
                ssh_keys=data["ssh"]["rootKeys"],
            )
        elif username == data["username"]:
            if "sshKeys" not in data:
                data["sshKeys"] = []

            return User(
                user_type=UserType.PRIMARY,
                username=username,
                ssh_keys=data["sshKeys"],
            )
        else:
            for user in data["users"]:
                if user["username"] == username:
                    if "sshKeys" not in user:
                        user["sshKeys"] = []

                    return User(
                        user_type=UserType.NORMAL,
                        username=username,
                        ssh_keys=user["sshKeys"],
                    )
        return None