# vim: ts=4:sw=4:expandtab
# -*- coding: UTF-8 -*-

# BleachBit
# Copyright (C) 2008-2025 Andrew Ziem
# https://www.bleachbit.org
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


"""
Wipe memory
"""

import logging
import os
import re
import subprocess
import sys

from bleachbit import FileUtilities
from bleachbit import General
from bleachbit.Language import get_text as _
from bleachbit.Wipe import wipe_contents

logger = logging.getLogger(__name__)


def count_swap_linux():
    """Count the number of swap devices in use"""
    count = 0
    with open("/proc/swaps") as f:
        for line in f:
            if line[0] == '/':
                count += 1
    return count


def get_proc_swaps():
    """Return the output of 'swapon -s'"""
    # Usually 'swapon -s' is identical to '/proc/swaps'
    # Here is one exception:
    # https://bugs.launchpad.net/ubuntu/+source/bleachbit/+bug/1092792
    (rc, stdout, _stderr) = General.run_external(['swapon', '-s'])
    if 0 == rc:
        return stdout
    logger.debug(
        _("The command 'swapoff -s' failed, so falling back to /proc/swaps for swap information."))
    return open("/proc/swaps").read()


def parse_swapoff(swapoff):
    """Parse the output of swapoff and return the device name"""
    # English is 'swapoff on /dev/sda5' but German is 'swapoff für ...'
    # Example output in English with LVM and hyphen: 'swapoff on /dev/mapper/lubuntu-swap_1'
    # This matches swap devices and swap files
    ret = re.search(r'^swapoff (\w* )?(/[\w/.-]+)$', swapoff)
    if not ret:
        # no matches
        return None
    return ret.group(2)


def disable_swap_linux():
    """Disable Linux swap and return list of devices"""
    if 0 == count_swap_linux():
        return
    logger.debug(_("Disabling swap."))
    args = ["swapoff", "-a", "-v"]
    (rc, stdout, stderr) = General.run_external(args)
    if 0 != rc:
        raise RuntimeError(stderr.replace("\n", ""))
    devices = []
    for line in stdout.split('\n'):
        line = line.replace('\n', '')
        if '' == line:
            continue
        ret = parse_swapoff(line)
        if ret is None:
            raise RuntimeError(f"Unexpected output:\nargs='{args}'\n"
                               f"stdout='{stdout}'\nstderr='{stderr}'")
        devices.append(ret)
    return devices


def enable_swap_linux():
    """Enable Linux swap"""
    logger.debug(_("Re-enabling swap."))
    args = ["swapon", "-a"]
    p = subprocess.Popen(args, stderr=subprocess.PIPE)
    p.wait()
    outputs = p.communicate()
    if 0 != p.returncode:
        raise RuntimeError(outputs[1].replace("\n", ""))


def make_self_oom_target_linux():
    """Make the current process the primary target for Linux out-of-memory killer"""
    path = f'/proc/{os.getpid()}/oom_score_adj'
    if os.path.exists(path):
        with open(path, 'w', encoding='utf-8') as f:
            f.write('1000')
    # OOM likes nice processes
    logger.debug(_("Setting nice value %d for this process."), os.nice(19))
    # OOM prefers non-privileged processes
    try:
        uid = General.get_real_uid()
        if uid > 0:
            # TRANSLATORS: Debug message when a process gives up root/admin privileges.
            # %(pid)d is the integer process ID; %(uid)d is the integer user ID to switch to.
            drop_msg = _("Dropping privileges of process ID %(pid)d to user ID %(uid)d.")
            logger.debug(drop_msg, {'pid': os.getpid(), 'uid': uid})
            os.seteuid(uid)
    except:
        logger.exception('Error when dropping privileges')


def fill_memory_linux():
    """Fill unallocated memory"""
    report_free()
    allocbytes = int(physical_free() * 0.4)
    if allocbytes < 1024:
        return
    bytes_str = FileUtilities.bytes_to_human(allocbytes)
    # TRANSLATORS: The variable is a quantity like 5kB
    logger.info(_("Allocating and wiping %s of memory."),
                bytes_str)
    try:
        buf = '\x00' * allocbytes
    except MemoryError:
        pass
    else:
        fill_memory_linux()
        # TRANSLATORS: The variable is a quantity like 5kB
        logger.debug(_("Freeing %s of memory."), bytes_str)
        del buf
    report_free()


def get_swap_size_linux(device, proc_swaps=None):
    """Return the size of the partition in bytes"""
    if proc_swaps is None:
        proc_swaps = get_proc_swaps()
    line = proc_swaps.split('\n')[0]
    if not re.search(r'Filename\s+Type\s+Size', line):
        raise RuntimeError("Unexpected first line in swap summary '%s'" % line)
    for line in proc_swaps.split('\n')[1:]:
        ret = re.search(r"%s\s+\w+\s+([0-9]+)\s" % device, line)
        if ret:
            return int(ret.group(1)) * 1024
    raise RuntimeError("error: cannot find size of swap device '%s'\n%s" %
                       (device, proc_swaps))


def get_swap_uuid(device):
    """Find the UUID for the swap device"""
    uuid = None
    args = ['blkid', device, '-s', 'UUID']
    (_rc, stdout, _stderr) = General.run_external(args)
    for line in stdout.split('\n'):
        # example: /dev/sda5: UUID="ee0e85f6-6e5c-42b9-902f-776531938bbf"
        ret = re.search(r"^%s: UUID=\"([a-z0-9-]+)\"" % device, line)
        if ret is not None:
            uuid = ret.group(1)
    # TRANSLATORS: Debug message. 'Found' is a past tense verb (short for
    # "Found [that] the UUID for swap device ..."). %(device)s is the device
    # path (e.g., /dev/sda5); %(uuid)s is a UUID string
    # (e.g., ee0e85f6-6e5c-42b9-902f-776531938bbf). Do not translate variables.
    logger.debug(_("Found UUID for swap device %(device)s is %(uuid)s."),
                 {'device': device, 'uuid': uuid})
    return uuid


def physical_free_darwin(run_vmstat=None):
    def parse_line(k, v):
        return k, int(v.strip(" ."))

    def get_page_size(line):
        m = re.match(
            r"Mach Virtual Memory Statistics: \(page size of (\d+) bytes\)",
            line)
        if m is None:
            raise RuntimeError("Can't parse vm_stat output")
        return int(m.groups()[0])
    if run_vmstat is None:
        def run_vmstat():
            return subprocess.check_output(["vm_stat"])
    output = iter(run_vmstat().split("\n"))
    page_size = get_page_size(next(output))
    vm_stat = dict(parse_line(*l.split(":")) for l in output if l != "")
    return vm_stat["Pages free"] * page_size


def physical_free_linux():
    """Return the physical free memory on Linux"""
    free_bytes = 0
    with open("/proc/meminfo") as f:
        for line in f:
            line = line.replace("\n", "")
            ret = re.search(r'(MemFree|Cached):[ ]*([0-9]*) kB', line)
            if ret is not None:
                kb = int(ret.group(2))
                free_bytes += kb * 1024
    if free_bytes > 0:
        return free_bytes
    else:
        raise Exception("unknown")


def physical_free_windows():
    """Return physical free memory on Windows"""

    from ctypes import c_ulong, c_ulonglong, Structure, sizeof, windll, byref

    class MEMORYSTATUSEX(Structure):
        _fields_ = [
            ('dwLength', c_ulong),
            ('dwMemoryLoad', c_ulong),
            ('ullTotalPhys', c_ulonglong),
            ('ullAvailPhys', c_ulonglong),
            ('ullTotalPageFile', c_ulonglong),
            ('ullAvailPageFile', c_ulonglong),
            ('ullTotalVirtual', c_ulonglong),
            ('ullAvailVirtual', c_ulonglong),
            ('ullExtendedVirtual', c_ulonglong),
        ]

    def GlobalMemoryStatusEx():
        x = MEMORYSTATUSEX()
        x.dwLength = sizeof(x)
        windll.kernel32.GlobalMemoryStatusEx(byref(x))
        return x

    z = GlobalMemoryStatusEx()
    return z.ullAvailPhys


def physical_free():
    if sys.platform == 'linux':
        return physical_free_linux()
    elif 'win32' == sys.platform:
        return physical_free_windows()
    elif 'darwin' == sys.platform:
        return physical_free_darwin()
    else:
        raise RuntimeError('unsupported platform for physical_free()')


def report_free():
    """Report free memory"""
    bytes_free = physical_free()
    bytes_str = FileUtilities.bytes_to_human(bytes_free)
    # TRANSLATORS: The variable is a quantity like 5kB
    logger.debug(_("Physical free memory is %s."),
                 bytes_str)


def wipe_swap_linux(devices, proc_swaps):
    """Shred the Linux swap file and then reinitialize it"""
    if devices is None:
        return
    if 0 < count_swap_linux():
        raise RuntimeError('Cannot wipe swap while it is in use')
    for device in devices:
        # if '/cryptswap' in device:
        #    logger.info('Skipping encrypted swap device %s.', device)
        #    continue
        # TRANSLATORS: The variable is a device like /dev/sda2
        logger.info(_("Wiping the swap device %s."), device)
        safety_limit_bytes = 29 * 1024 ** 3  # 29 gibibytes
        actual_size_bytes = get_swap_size_linux(device, proc_swaps)
        if actual_size_bytes > safety_limit_bytes:
            raise RuntimeError(
                f'swap device {device} is larger ({actual_size_bytes})'
                f' than expected ({safety_limit_bytes})')
        uuid = get_swap_uuid(device)
        # wipe
        wipe_contents(device, truncate=False)
        # reinitialize
        # TRANSLATORS: The variable is a device like /dev/sda2
        logger.debug(_("Reinitializing the swap device %s."), device)
        args = ['mkswap', device]
        if uuid:
            args.append("-U")
            args.append(uuid)
        (rc, _stdout, stderr) = General.run_external(args)
        if 0 != rc:
            raise RuntimeError(stderr.replace("\n", ""))


def wipe_memory():
    """Wipe unallocated memory"""
    for cmd in ('swapon', 'swapoff', 'blkid'):
        if not FileUtilities.exe_exists(cmd):
            raise RuntimeError(f"wipe_memory: Command {cmd} not found")
    # cache the file because 'swapoff' changes it
    proc_swaps = get_proc_swaps()
    devices = disable_swap_linux()
    yield True  # process GTK+ idle loop
    # TRANSLATORS: The variable is a device like /dev/sda2
    logger.debug(_("Detected these swap devices: %s"), str(devices))
    wipe_swap_linux(devices, proc_swaps)
    yield True
    child_pid = os.fork()
    if 0 == child_pid:
        make_self_oom_target_linux()
        fill_memory_linux()
        os._exit(0)
    else:
        # TRANSLATORS: This is a debugging message that the parent process
        # is waiting for the child process. %(parent_pid)d is the parent
        # process ID; %(child_pid)d is the child process ID.
        logger.debug(_("The function wipe_memory() with process ID %(parent_pid)d is "
                       "waiting for child process ID %(child_pid)d."),
                     {'parent_pid': os.getpid(), 'child_pid': child_pid})
        rc = os.waitpid(child_pid, 0)[1]
        if rc not in [0, 9]:
            logger.warning(
                _("The child memory-wiping process returned code %d."), rc)
    enable_swap_linux()
    yield 0  # how much disk space was recovered
