dotfiles/files/.local/bin/motd

#!/usr/bin/env python3

import math
import shutil
import sys
import os
import re

from subprocess import check_call, DEVNULL, CalledProcessError

try:
    import psutil
except:
    pass

try:
    from pyfiglet import figlet_format
except ImportError:
    def figlet_format(string, *args, **kwargs):
       return string + '\n'

term_r, term_c = os.popen('stty size', 'r').read().split()

class Termutils:
    BLACK = '\33[30m'
    RED = '\33[31m'
    GREEN = '\33[32m'
    YELLOW = '\33[33m'
    BLUE = '\33[34m'
    MAGENTA = '\33[35m'
    CYAN = '\33[36m'
    WHITE = '\33[37m'
    GREY = '\33[30;1m'
    DISABLE = '\033[0m'

    HOSTNAME_COLOR = '\033[38;5;%sm' % os.environ.get('ZSH_THEME_HOSTNAME_COLOR', '004')

    PADDING = '  '

    @staticmethod
    def print_padded(*arg, **kwarg):
        print(Termutils.PADDING, end='')
        print(*arg, **kwarg)

    @staticmethod
    def print_progress_bar(percent, label='', length=60):
        percent_color = Termutils.GREEN
        bar = '['
        for i in range(0, length):
            if i >= math.ceil(percent * length):
                bar += Termutils.GREY
            elif i == 0:
                bar += Termutils.GREEN
            elif i == math.ceil(length * 0.6):
                bar += Termutils.YELLOW
                percent_color = Termutils.YELLOW
            elif i == math.ceil(length * 0.9):
                bar += Termutils.RED
                percent_color = Termutils.RED

            bar += '='
        bar += '{}]{}{:7.2%}{} {}'.format(
            Termutils.DISABLE,
            percent_color,
            percent,
            Termutils.DISABLE,
            label
        )
        Termutils.print_padded(bar)

def to_gb(bytes_count):
    return bytes_count/1024/1024/1024

def print_hostname():
    print(
          Termutils.HOSTNAME_COLOR
        + figlet_format(os.uname()[1], 'slant')
        + Termutils.DISABLE
        + '\n'
        + os.popen('uptime').read().strip()
        + '\n'
    )

def print_services():
    if 'MOTD_SERVICES' in os.environ:
        print('Services:')
        services = list(set(os.environ['MOTD_SERVICES'].split()))
        services.sort()
        for service in services:
            Termutils.print_padded('{:<15.15}'.format(service), end='')
            try:
                check_call(['systemctl', 'is-active', service], stdout=DEVNULL, stderr=DEVNULL)
                print('{}  ● active{}'.format(Termutils.GREEN, Termutils.DISABLE))
            except CalledProcessError:
                print('{}  ● inactive{}'.format(Termutils.RED, Termutils.DISABLE))
        print('')


def print_mem():
    try:
        memory = psutil.virtual_memory()
        print('{:<18}  {:6.2f}G used  {:6.2f}G free  {:6.2f}G total'.format(
            'Memory:',
            to_gb(memory.used),
            to_gb(memory.free),
            to_gb(memory.total)
        ))
        Termutils.print_progress_bar(memory.used / memory.total)
        print('')
    except NameError:
        pass


def print_fs():
    print('Filesystems:')
    with open('/etc/mtab') as mtab:
        for line in mtab:
            line = line.strip()
            if line.startswith('/'):
                mount_point = line.split()[1]
                usage = shutil.disk_usage(mount_point)
                Termutils.print_padded('{:<15.15}  {:7.2f}G used  {:7.2f}G free  {:7.2f}G total'.format(
                    mount_point,
                    to_gb(usage.used),
                    to_gb(usage.free),
                    to_gb(usage.total)
                ))
                Termutils.print_progress_bar(
                    percent=(usage.used / usage.total)
                )
    print('')

def print_tmux():
    try:
        check_call(['tmux', 'info'], stdout=DEVNULL, stderr=DEVNULL)
        print('tmux:')
        tmux_session = os.popen('tmux display-message -p "#S"', 'r').read().strip()
        with os.popen('tmux list-sessions', 'r') as tmux_ls:
            for line in tmux_ls:
                line = line.strip()
                if line.startswith(tmux_session) and 'TMUX' in os.environ:
                    print(Termutils.GREEN, end='')
                Termutils.print_padded(
                      line
                    + Termutils.DISABLE)
        print('')
    except CalledProcessError:
        pass


print_hostname()
print_services()
print_mem()
print_fs()
print_tmux()