#!/usr/bin/python

# Copyright (c) 2018, Emmanouil Kampitakis <info@kampitakis.de>
# Copyright (c) 2018, William Leemans <willie@elaba.net>

# GNU General Public License v3.0+ (see LICENSES/GPL-3.0-or-later.txt or https://www.gnu.org/licenses/gpl-3.0.txt)
# SPDX-License-Identifier: GPL-3.0-or-later

from __future__ import annotations

DOCUMENTATION = r"""
module: xfs_quota
short_description: Manage quotas on XFS filesystems
description:
  - Configure quotas on XFS filesystems.
  - Before using this module /etc/projects and /etc/projid need to be configured.
author:
  - William Leemans (@bushvin)
extends_documentation_fragment:
  - community.general._attributes
attributes:
  check_mode:
    support: full
  diff_mode:
    support: none
options:
  type:
    description:
      - The XFS quota type.
    type: str
    required: true
    choices:
      - user
      - group
      - project
  name:
    description:
      - The name of the user, group or project to apply the quota to, if other than default.
    type: str
  mountpoint:
    description:
      - The mount point on which to apply the quotas.
    type: str
    required: true
  bhard:
    description:
      - Hard blocks quota limit.
      - This argument supports human readable sizes.
    type: str
  bsoft:
    description:
      - Soft blocks quota limit.
      - This argument supports human readable sizes.
    type: str
  ihard:
    description:
      - Hard inodes quota limit.
    type: int
  isoft:
    description:
      - Soft inodes quota limit.
    type: int
  rtbhard:
    description:
      - Hard realtime blocks quota limit.
      - This argument supports human readable sizes.
    type: str
  rtbsoft:
    description:
      - Soft realtime blocks quota limit.
      - This argument supports human readable sizes.
    type: str
  state:
    description:
      - Whether to apply the limits or remove them.
      - When removing limit, they are set to 0, and not quite removed.
    type: str
    default: present
    choices:
      - present
      - absent

requirements:
  - xfsprogs
"""

EXAMPLES = r"""
- name: Set default project soft and hard limit on /opt of 1g
  community.general.xfs_quota:
    type: project
    mountpoint: /opt
    bsoft: 1g
    bhard: 1g
    state: present

- name: Remove the default limits on /opt
  community.general.xfs_quota:
    type: project
    mountpoint: /opt
    state: absent

- name: Set default soft user inode limits on /home of 1024 inodes and hard of 2048
  community.general.xfs_quota:
    type: user
    mountpoint: /home
    isoft: 1024
    ihard: 2048
"""

RETURN = r"""
bhard:
  description: The current C(bhard) setting in bytes.
  returned: always
  type: int
  sample: 1024
bsoft:
  description: The current C(bsoft) setting in bytes.
  returned: always
  type: int
  sample: 1024
ihard:
  description: The current C(ihard) setting in bytes.
  returned: always
  type: int
  sample: 100
isoft:
  description: The current C(isoft) setting in bytes.
  returned: always
  type: int
  sample: 100
rtbhard:
  description: The current C(rtbhard) setting in bytes.
  returned: always
  type: int
  sample: 1024
rtbsoft:
  description: The current C(rtbsoft) setting in bytes.
  returned: always
  type: int
  sample: 1024
"""

import grp
import os
import pwd

from ansible.module_utils.basic import AnsibleModule, human_to_bytes


def main():
    module = AnsibleModule(
        argument_spec=dict(
            bhard=dict(type="str"),
            bsoft=dict(type="str"),
            ihard=dict(type="int"),
            isoft=dict(type="int"),
            mountpoint=dict(type="str", required=True),
            name=dict(type="str"),
            rtbhard=dict(type="str"),
            rtbsoft=dict(type="str"),
            state=dict(type="str", default="present", choices=["absent", "present"]),
            type=dict(type="str", required=True, choices=["group", "project", "user"]),
        ),
        supports_check_mode=True,
    )

    quota_type = module.params["type"]
    name = module.params["name"]
    mountpoint = module.params["mountpoint"]
    bhard = module.params["bhard"]
    bsoft = module.params["bsoft"]
    ihard = module.params["ihard"]
    isoft = module.params["isoft"]
    rtbhard = module.params["rtbhard"]
    rtbsoft = module.params["rtbsoft"]
    state = module.params["state"]

    xfs_quota_bin = module.get_bin_path("xfs_quota", True)

    if bhard is not None:
        bhard = human_to_bytes(bhard)

    if bsoft is not None:
        bsoft = human_to_bytes(bsoft)

    if rtbhard is not None:
        rtbhard = human_to_bytes(rtbhard)

    if rtbsoft is not None:
        rtbsoft = human_to_bytes(rtbsoft)

    result = dict(
        changed=False,
    )

    if not os.path.ismount(mountpoint):
        module.fail_json(msg=f"Path '{mountpoint}' is not a mount point", **result)

    mp = get_fs_by_mountpoint(mountpoint)
    if mp is None:
        module.fail_json(
            msg=f"Path '{mountpoint}' is not a mount point or not located on an xfs file system.", **result
        )

    if quota_type == "user":
        type_arg = "-u"
        quota_default = "root"
        if name is None:
            name = quota_default

        if (
            "uquota" not in mp["mntopts"]
            and "usrquota" not in mp["mntopts"]
            and "quota" not in mp["mntopts"]
            and "uqnoenforce" not in mp["mntopts"]
            and "qnoenforce" not in mp["mntopts"]
        ):
            module.fail_json(
                msg=f"Path '{mountpoint}' is not mounted with the uquota/usrquota/quota/uqnoenforce/qnoenforce option.",
                **result,
            )
        try:
            pwd.getpwnam(name)
        except KeyError:
            module.fail_json(msg=f"User '{name}' does not exist.", **result)

    elif quota_type == "group":
        type_arg = "-g"
        quota_default = "root"
        if name is None:
            name = quota_default

        if "gquota" not in mp["mntopts"] and "grpquota" not in mp["mntopts"] and "gqnoenforce" not in mp["mntopts"]:
            module.fail_json(
                msg=f"Path '{mountpoint}' is not mounted with the gquota/grpquota/gqnoenforce option. (current options: {mp['mntopts']})",
                **result,
            )
        try:
            grp.getgrnam(name)
        except KeyError:
            module.fail_json(msg=f"User '{name}' does not exist.", **result)

    elif quota_type == "project":
        type_arg = "-p"
        quota_default = "#0"
        if name is None:
            name = quota_default

        if "pquota" not in mp["mntopts"] and "prjquota" not in mp["mntopts"] and "pqnoenforce" not in mp["mntopts"]:
            module.fail_json(
                msg=f"Path '{mountpoint}' is not mounted with the pquota/prjquota/pqnoenforce option.", **result
            )

        if name != quota_default and not os.path.isfile("/etc/projects"):
            module.fail_json(msg="Path '/etc/projects' does not exist.", **result)

        if name != quota_default and not os.path.isfile("/etc/projid"):
            module.fail_json(msg="Path '/etc/projid' does not exist.", **result)

        if name != quota_default and name is not None and get_project_id(name) is None:
            module.fail_json(msg=f"Entry '{name}' has not been defined in /etc/projid.", **result)

        prj_set = True
        if name != quota_default:
            cmd = f"project {name}"
            rc, stdout, stderr = exec_quota(module, xfs_quota_bin, cmd, mountpoint)
            if rc != 0:
                result["cmd"] = cmd
                result["rc"] = rc
                result["stdout"] = stdout
                result["stderr"] = stderr
                module.fail_json(msg="Could not get project state.", **result)
            else:
                for line in stdout.split("\n"):
                    if "Project Id '%s' - is not set." in line or "project identifier is not set" in line:
                        prj_set = False
                        break

        if state == "present" and not prj_set:
            if not module.check_mode:
                cmd = f"project -s {name}"
                rc, stdout, stderr = exec_quota(module, xfs_quota_bin, cmd, mountpoint)
                if rc != 0:
                    result["cmd"] = cmd
                    result["rc"] = rc
                    result["stdout"] = stdout
                    result["stderr"] = stderr
                    module.fail_json(msg="Could not get quota realtime block report.", **result)

            result["changed"] = True

        elif state == "absent" and prj_set and name != quota_default:
            if not module.check_mode:
                cmd = f"project -C {name}"
                rc, stdout, stderr = exec_quota(module, xfs_quota_bin, cmd, mountpoint)
                if rc != 0:
                    result["cmd"] = cmd
                    result["rc"] = rc
                    result["stdout"] = stdout
                    result["stderr"] = stderr
                    module.fail_json(msg="Failed to clear managed tree from project quota control.", **result)

            result["changed"] = True

    current_bsoft, current_bhard = quota_report(module, xfs_quota_bin, mountpoint, name, quota_type, "b")
    current_isoft, current_ihard = quota_report(module, xfs_quota_bin, mountpoint, name, quota_type, "i")
    current_rtbsoft, current_rtbhard = quota_report(module, xfs_quota_bin, mountpoint, name, quota_type, "rtb")

    # Set limits
    if state == "absent":
        bhard = 0
        bsoft = 0
        ihard = 0
        isoft = 0
        rtbhard = 0
        rtbsoft = 0

        # Ensure that a non-existing quota does not trigger a change
        current_bsoft = current_bsoft if current_bsoft is not None else 0
        current_bhard = current_bhard if current_bhard is not None else 0
        current_isoft = current_isoft if current_isoft is not None else 0
        current_ihard = current_ihard if current_ihard is not None else 0
        current_rtbsoft = current_rtbsoft if current_rtbsoft is not None else 0
        current_rtbhard = current_rtbhard if current_rtbhard is not None else 0

    result["xfs_quota"] = dict(
        bsoft=current_bsoft,
        bhard=current_bhard,
        isoft=current_isoft,
        ihard=current_ihard,
        rtbsoft=current_rtbsoft,
        rtbhard=current_rtbhard,
    )

    limit = []
    if bsoft is not None and int(bsoft) != current_bsoft:
        limit.append(f"bsoft={bsoft}")
        result["bsoft"] = int(bsoft)

    if bhard is not None and int(bhard) != current_bhard:
        limit.append(f"bhard={bhard}")
        result["bhard"] = int(bhard)

    if isoft is not None and isoft != current_isoft:
        limit.append(f"isoft={isoft}")
        result["isoft"] = isoft

    if ihard is not None and ihard != current_ihard:
        limit.append(f"ihard={ihard}")
        result["ihard"] = ihard

    if rtbsoft is not None and int(rtbsoft) != current_rtbsoft:
        limit.append(f"rtbsoft={rtbsoft}")
        result["rtbsoft"] = int(rtbsoft)

    if rtbhard is not None and int(rtbhard) != current_rtbhard:
        limit.append(f"rtbhard={rtbhard}")
        result["rtbhard"] = int(rtbhard)

    if len(limit) > 0:
        if not module.check_mode:
            if name == quota_default:
                cmd = f"limit {type_arg} -d {' '.join(limit)}"
            else:
                cmd = f"limit {type_arg} {' '.join(limit)} {name}"

            rc, stdout, stderr = exec_quota(module, xfs_quota_bin, cmd, mountpoint)
            if rc != 0:
                result["cmd"] = cmd
                result["rc"] = rc
                result["stdout"] = stdout
                result["stderr"] = stderr
                module.fail_json(msg="Could not set limits.", **result)

        result["changed"] = True

    module.exit_json(**result)


def quota_report(module, xfs_quota_bin, mountpoint, name, quota_type, used_type):
    soft = None
    hard = None

    if quota_type == "project":
        type_arg = "-p"
    elif quota_type == "user":
        type_arg = "-u"
    elif quota_type == "group":
        type_arg = "-g"

    if used_type == "b":
        used_arg = "-b"
        used_name = "blocks"
        factor = 1024
    elif used_type == "i":
        used_arg = "-i"
        used_name = "inodes"
        factor = 1
    elif used_type == "rtb":
        used_arg = "-r"
        used_name = "realtime blocks"
        factor = 1024

    rc, stdout, stderr = exec_quota(module, xfs_quota_bin, f"report {type_arg} {used_arg}", mountpoint)

    if rc != 0:
        result = dict(
            changed=False,
            rc=rc,
            stdout=stdout,
            stderr=stderr,
        )
        module.fail_json(msg=f"Could not get quota report for {used_name}.", **result)

    for line in stdout.split("\n"):
        line = line.strip().split()
        if len(line) > 3 and line[0] == name:
            soft = int(line[2]) * factor
            hard = int(line[3]) * factor
            break

    return soft, hard


def exec_quota(module, xfs_quota_bin, cmd, mountpoint):
    cmd = [xfs_quota_bin, "-x", "-c", cmd, mountpoint]
    (rc, stdout, stderr) = module.run_command(cmd)
    if (
        "XFS_GETQUOTA: Operation not permitted" in stderr.split("\n")
        or rc == 1
        and "xfs_quota: cannot set limits: Operation not permitted" in stderr.split("\n")
    ):
        module.fail_json(msg="You need to be root or have CAP_SYS_ADMIN capability to perform this operation")

    return rc, stdout, stderr


def get_fs_by_mountpoint(mountpoint):
    mpr = None
    with open("/proc/mounts") as s:
        for line in s.readlines():
            mp = line.strip().split()
            if len(mp) == 6 and mp[1] == mountpoint and mp[2] == "xfs":
                mpr = dict(zip(["spec", "file", "vfstype", "mntopts", "freq", "passno"], mp))
                mpr["mntopts"] = mpr["mntopts"].split(",")
                break
    return mpr


def get_project_id(name):
    prjid = None
    with open("/etc/projid") as s:
        for line in s.readlines():
            line = line.strip().partition(":")
            if line[0] == name:
                prjid = line[2]
                break

    return prjid


if __name__ == "__main__":
    main()
