#!/usr/bin/python3
# -*- coding: utf-8 -*-

import os
import subprocess

message_printed = False

def print_message():
    # Print user-friendly message
    global message_printed
    if not message_printed:
        print()
        print('Bylo detekováno zvětšení virtuálního disku.')
        print('Systém nyní upraví nastavení diskových oddílů.')
        print()
        message_printed = True

def get_sizes():
    # Returns device sizes in MB
    sizes = {}
    lsblk = subprocess.run(['/bin/lsblk', '-bnro', 'PATH,SIZE'], capture_output=True, check=True).stdout.decode()
    for line in lsblk.splitlines():
        line = line.split()
        sizes[line[0]] = int(line[1])//1024**2
    return sizes

# Determine partition and disk paths
with open('/etc/crypttab') as f:
    crypttab = dict(line.split()[:2] for line in f.read().splitlines())
partition = subprocess.run(['/sbin/findfs', crypttab['system']], capture_output=True, check=True).stdout.decode().strip()
disk = partition.rstrip('0123456789')

# Resize physical partition when there is more than 10M unallocated
sizes = get_sizes()
free_space = sizes[disk] - sum([size for part,size in sizes.items() if part.startswith(disk) and part[len(disk):].isdigit()])
if free_space > 10:
    print_message()
    try:
        # Force busybox fdisk as util-linux fdisk breaks the subsequent partx command by notifying kernel too hard
        input = 'd\n2\nn\np\n2\n\n\nt\n2\n8e\nw'.encode()
        subprocess.run(['/bin/busybox', 'fdisk', disk], input=input, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
    except subprocess.CalledProcessError as e:
        # Returns code 1  because of the nofication failure, but the resize is successful anyway
        if e.returncode != 1:
            raise
    # Re-read partition (partx is from util-linux, issues BLKPG_RESIZE_PARTITION ioctl to update the kernel partition table)
    subprocess.run(['/usr/sbin/partx', '-u', partition], stdout=subprocess.DEVNULL, check=True)
    sizes = get_sizes()

# Resize dmcrypt when there is more than 10M of free space (16M is consumed by dmcrypt)
free_space = sizes[partition] - sizes['/dev/mapper/system'] - 16
if free_space > 10:
    print_message()
    # Create locking dir for cryptsetup to suppress warning
    try:
        os.mkdir('/run/cryptsetup', 0o700)
    except FileExistsError:
        pass
    # Ask for password 3 times, then fail
    fails = 0
    while True:
        try:
            subprocess.run(['/sbin/cryptsetup', 'resize', 'system'], check=True)
            break
        except (subprocess.CalledProcessError, KeyboardInterrupt):
            fails += 1
            if fails == 3:
                raise
    sizes = get_sizes()

# Resize LVM PV when there is more than 10M of free space
free_space = sizes['/dev/mapper/system'] - sum(size for part,size in sizes.items() if part.startswith('/dev/mapper/vg0'))
if free_space > 10:
    print_message()
    subprocess.run(['/sbin/pvresize', '/dev/mapper/system'], stdout=subprocess.DEVNULL, check=True)
# Create swap if it doesn't exist yet and there's at least 4 GB of free space
if '/dev/mapper/vg0-swap' not in sizes and free_space > 4*1024:
    subprocess.run(['/sbin/lvcreate', '-L', '4G', 'vg0', '-n', 'swap'], stdout=subprocess.DEVNULL, check=True)
    subprocess.run(['/sbin/mkswap', '/dev/mapper/vg0-swap'], stdout=subprocess.DEVNULL, check=True)
    subprocess.run(['/sbin/swapon', '/dev/mapper/vg0-swap'], stdout=subprocess.DEVNULL, check=True)
    sizes = get_sizes()
    free_space = sizes['/dev/mapper/system'] - sum(size for part,size in sizes.items() if part.startswith('/dev/mapper/vg0'))
# Resize LVM LV when there is more than 10M of free space
if free_space > 10:
    subprocess.run(['/sbin/lvextend', '-l', '+100%FREE', '/dev/mapper/vg0-root'], stdout=subprocess.DEVNULL, check=True)
    subprocess.run(['/usr/sbin/resize2fs', '/dev/mapper/vg0-root'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)