archeyX

#!/usr/bin/env python3
"""
/usr/local/bin/archeyX

Description:
  Archey X is a simple system information tool written in Python.

Licence:
  Distributed under GPL-3.0-or-later (GNU General Public License v3 or later).
  See LICENSE, or <https://www.gnu.org/licenses/gpl-3.0.html>.

Author notes:
  This is a forked/ported version of the original script: Archey, by Melik Manukyan.
  It's made to be cross platform, and able to auto detect different distros.

Original author(s) and copyright:
  (c) 2010 Melik Manukyan <melik@archlinux.us>
  (c) 2010 David Vazgenovich Shakaryan <dvshakaryan@gmail.com>

  ASCII art by Brett Bohnenkamper <kittykatt@silverirc.com>
  Changes Jerome Launay <jerome@projet-libre.org>
  Fedora support by YeOK <yeok@henpen.org>
"""

# Import libraries/resources
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from datetime import datetime, timedelta
from distutils.spawn import find_executable as fExec
from getpass import getuser
from itertools import chain
from math import ceil
from os import getenv, uname
from platform import freedesktop_os_release as os_release
from re import sub
from socket import gethostname
from subprocess import PIPE, Popen
from time import time
import sys


# @todo: Implement some sort of raise error for versioncheck
# It will bump to 3.10 soon.
# Inform about python3.10+ (force an SyntaxError)
INFO = f'Archey X requires Python 3.6+ - Please upgrade...'

# Version: use --help and --version
VERSION = 'v0.99-20240621'
COPYRIGHT = f'Copyright (c) {datetime.now().year} Eric F'
GPL_INFO = '''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.\n\nYou should have received a copy of the GNU General Public License \
along\nwith this program. If not, see <https://www.gnu.org/licenses/>.'''

# Settings
DISK_WARNING: int = 1                # default: 1 (set 0 to disable)
DISK_WARNING_THRESHOLD: int = 20     # default: 20 (% of diskspace left)


# Dictionaries, lists and tuples
# --------------------------------
COLOR_PAIRS: tuple = (
    (33, 97), (34, 34), (37, 30), (36, 36),  # yellow/white, blue/blue, grey/grey, cyan/cyan
    (34, 97), (31, 31), (97, 35), (97, 32),  # blue/white, red/red, white/purple, white/green
    (37, 32), (32, 30), (32, 32), (34, 32),  # grey/green, green/grey, green/green, blue/green
)

COLOR_VARIES: tuple = ('\x1b[0;31m', '\x1b[0;32m', '\x1b[0;33m', '\x1b[5;31m', '\x1b[0m')
COLOR_DICT: dict = {f'color{c:02d}': (f'\x1b[0;{n}m', f'\x1b[1;{b}m')
                    for c, (n, b) in enumerate(COLOR_PAIRS, 1)}

DIST_DICT: dict = {
    'GNU/Linux': 'color01', 'Arch': 'color02', 'Crunchbang': 'color03',
    'Darwin': 'color04', 'Mac OS X': 'color04', 'Fedora': 'color05',
    'Alpine': 'color05', 'Antergos': 'color05', 'elementary OS': 'color05',
    'RHEL': 'color06', 'Debian': 'color06', 'Ubuntu': 'color06', 'Gentoo': 'color07',
    'Linux Mint': 'color08', 'SLES': 'color09', 'Manjaro': 'color09',
    'OpenSUSE': 'color11', 'Rocky': 'color11', 'Alma': 'color12', 'CentOS': 'color12',
}

DE_DICT: dict = {
    'cinnamon-sessio': 'Cinnamon', 'Finder': 'Finder', 'gnome-session': 'GNOME',
    'ksmserver': 'KDE', 'lxsession': 'LXDE', 'mate-session': 'MATE',
    'mate-terminal': 'MATE', 'xfce4-session': 'Xfce'
}

WM_DICT: dict = {
    'awesome': 'Awesome', 'beryl': 'Beryl', 'blackbox': 'Blackbox', 'bspwm': 'Bspwm',
    'caja': 'Caja', 'compiz': 'Compiz', 'dwm': 'DWM', 'enlightenment': 'Enlightenment',
    'fluxbox': 'Fluxbox', 'fvwm': 'FVWM', 'herbstluftwm': 'Herbstluftwm', 'i3': 'i3',
    'icewm': 'IceWM', 'kwin': 'KWin', 'nautilus': 'Nautilus', 'metacity': 'Metacity',
    'musca': 'Musca', 'nemo': 'Nemo', 'openbox': 'Openbox', 'pekwm': 'PekWM',
    'ratpoison': 'Ratpoison', 'scrotwm': 'ScrotWM', 'wmaker': 'Window Maker',
    'wmfs': 'Wmfs', 'wmii': 'Wmii', 'xfwm4': 'Xfwm', 'xmonad': 'Xmonad',

    # OS X/macOS
    'Amethyst': 'Amethyst', 'AppGrid': 'AppGrid', 'Autumn': 'Autumn', 'Cinch': 'Cinch',
    'chunkwm': 'Chunkwm', 'Factory': 'Factory', 'Hammerspoon': 'Hammerspoon',
    'Mizage': 'Mizage', 'Moom': 'Moom', 'Optimal Layout': 'Optimal Layout',
    'Phoenix': 'Phoenix', 'ShiftIt': 'ShiftIt', 'SizeUp': 'SizeUp', 'Slate': 'Slate',
    'Slicer': 'Slicer', 'Spectacle': 'Spectacle', 'TylerWM': 'TylerWM', 'Xmonad': 'Xmonad'
}

MACOS_NAMES: tuple = (
    'Cheetah', 'Puma', 'Jaguar', 'Panther', 'Tiger', 'Leopard', 'Snow Leopard',
    'Lion', 'Mountain Lion', 'Mavericks', 'Yosemite', 'El Capitan', 'Sierra',
    'High Sierra', 'Mojave', 'Catalina', 'Big Sur', 'Monterey', 'Ventura',
    'Sonoma', 'Sequoia')

MACOS10: dict = {f'10.{x}': y for x, y in enumerate(MACOS_NAMES[:17])}
MACOS11: dict = {f'{11+x}': y for x, y in enumerate(MACOS_NAMES[16:])}

MACOS_DICT: dict = {**MACOS10, **MACOS11}

PM_DICT: dict = {
    'pacman': ['-Q'], 'dnf': ['rq', '--installed', '-Cq'], 'rpm': ['-qa'],
    'apk': ['info'], 'dpkg': ['--get-selections'], 'emerge': ['-ep', 'world'],
    'zypper': ['search', '--installed-only'],

    'port': ['echo', 'installed'], 'fink': ['list', '-i'], 'brew': ['list']     # OS X/macOS
}

LABELS: tuple = (
    'User', 'Hostname', 'Distro', 'Kernel', 'Uptime', 'WindowManager',
    'DesktopEnvironment', 'Shell', 'Terminal', 'Packages', 'CPU', 'RAM', 'Disk')

CLASSES: dict = {x: x for x in LABELS}

# List of filesystemes to be included in 'df'
FS_TYPES: tuple = (
    'apfs', 'btrfs', 'ext2', 'ext3', 'ext4', 'fat32', 'f2fs', 'fuseblk',
    'hfs', 'jfs', 'lxfs', 'msdos', 'nilfs2', 'ntfs', 'reiserfs', 'simfs',
    'tempfs', 'xfs', 'zfs')


# ArgumentParser - help/options
# --------------------------------
__version__ = f'%(prog)s: {VERSION}, {COPYRIGHT}\n\n{GPL_INFO}\v'
PARSER = ArgumentParser(prog='archeyX', formatter_class=RawDescriptionHelpFormatter,
                        add_help=False, conflict_handler='resolve')
PARSER.add_argument('-h', '--help', action='help', help='Show this help message and exit')
PARSER.add_argument('-V', '--version', action='version', version=__version__,
                    help='Show program\'s version number and exit')
PARSER.add_argument('-b', dest='brand', default=None, nargs=1, type=str)  # @todo: column width
PARSER.add_argument('--brand', dest='brand', default=None, nargs=1, type=str,
                    help='''You can pass in a Brand/distro name for another color.
                    The name must exist in DIST_DICT (ie. Arch Linux, Fedora, Debian etc.)''')

MACOS = PARSER.add_argument_group('OS X/macOS (only)')
MACOS.add_argument('-a', '--apple', dest='apple', action='store_true',
                   help='''If your system is Darwin you can pass `-a' to
                   display the Apple instead. (A left-over from the bash-script.)''')

if __name__ == "__main__":
    BRAND = None
    ARGS = PARSER.parse_args()
    if ARGS.brand and ARGS.brand[0] in DIST_DICT:
        BRAND = ARGS.brand[0]
    elif ARGS.apple is True:
        BRAND = 'Apple'


# Classes
# --------------------------------
class Utils():
    """Misc functions/tools"""

    @classmethod
    def x_popen(cls, cmd: list) -> str:
        """x_popen - less code and easier to use"""
        return Popen(cmd, stdout=PIPE).communicate()[0].decode('utf-8').rstrip('\n')

    def is_platform(self, platform: str = 'linux') -> bool:
        """is_platform - Linux, Darwin, etc"""
        return platform.lower() in sys.platform

    def x_process(self) -> list:
        """x_process - get the processes to find name of WM/DE"""
        if self.is_platform('Linux'):
            procs = self.x_popen(['ps', '-u', getuser(), '-o', 'comm', '--no-headers']).split('\n')
        elif self.is_platform('Darwin'):
            procs = self.x_popen(['ps', '-cu', getuser(), '-o', 'comm']).split('\n')
        return procs if procs else None


class Output(Utils):
    """Function and tools for the Output"""
    results = []
    results.extend([''] * (18 - len(LABELS)))

    def __init__(self) -> None:
        """default to GNU/Linux if no distro"""
        distro = self.x_distro().rstrip()
        self.distro = 'GNU/Linux' if distro not in DIST_DICT else distro

        brand_color = self.distro if BRAND not in DIST_DICT else BRAND
        self.color = DIST_DICT.get(brand_color, distro)

    def x_distro(self, name: str = 'name') -> str:
        """x_distro: determing what distro is being used, both name and name"""
        distro_name = distro_full = '' if self.is_platform('Darwin') else 'GNU/Linux'

        if self.is_platform('Darwin'):
            # @todo:: Check platform.mac_os(), to move away from popen
            get_name = self.x_popen(['sw_vers', '-productName'])
            get_version = _gv = self.x_popen(['sw_vers', '-productVersion'])
            osx_name = self.x_popen(['sw_vers', '-buildVersion'])

            get_name = 'macOS' if float(get_version[:-2]) >= float('10.12') else get_name \
                       or get_name if float(get_version[:-2]) <= float('10.7') else 'OS X'

            # 10.7, 10.8, etc - and then: 11, 12, 13, etc
            get_version = _gv[:2] if int(_gv[:2]) > 10 else _gv[:-2]

            # Get the fancy name (ie Tiger, Leopard, Lion, etc)
            # If not in the dict, display -buildVersion instead
            osx_name = MACOS_DICT.get(get_version, osx_name)

            distro_name, distro_full = 'Darwin', f'{get_name} {get_version} ({osx_name})'

        # (*BSD... Hopefully at some point)
        # elif self.is_platform('bsdD'):
        #     distro = distro_full = ''

        elif self.is_platform('Linux'):
            distro_name = distro_name = sub('Linux', '', os_release()['NAME'])
            distro_full = os_release()['PRETTY_NAME']

            distro_full = sub('Red Hat', 'RHEL', distro_full)
            distro_full = sub('S[Uu]SE', 'SLES', distro_full)

        x_distro_names = {'name': distro_name, 'pretty_name': distro_full}
        x_sub = sub('release | GNU/Linux|Enterprise |Server ', '', x_distro_names[name])

        return x_distro_names[name] if distro_name == 'Linux Mint' else x_sub

    def append(self, display: str) -> str:
        """append: preparing the output"""
        self.results.append(f'{COLOR_DICT[self.color][1]}{display.key}: '
                            f'{COLOR_VARIES[-1]}{display.value}')

    def output(self):
        """output: matching name/colors"""
        return LOGOS_DICT[self.distro].format(color=COLOR_DICT[self.color], results=self.results)


class User(Utils):
    """Get the username"""
    def __init__(self) -> None:
        self.key, self.value = self.__class__.__name__, getuser()


class Hostname(Utils):
    """Get the hostname"""
    def __init__(self) -> None:
        self.key, self.value = self.__class__.__name__, gethostname()


class Distro(Output, Utils):
    """Display the Distro name"""
    def __init__(self) -> None:
        self.key, self.value = self.__class__.__name__, self.x_distro('pretty_name')


class Kernel(Utils):
    """Get the kernel name"""
    def __init__(self) -> None:
        uname_srm = xnu = f'{uname()[0]} {uname()[2]} {uname()[4]}' # Ex: Linux 5.x.x x86_64
        #xnu = sub('.*:+', '', uname()[3]).strip('REL.*')           # xnu-123...
        kernel = xnu if self.is_platform('Darwin') else uname_srm

        self.key, self.value = self.__class__.__name__, kernel


class Uptime(Utils):
    """Uptime: N day(s), HH:MM"""
    def __init__(self) -> None:
        if self.is_platform('Darwin'):
            boot_time = self.x_popen(['sysctl', '-n', 'kern.boottime'])
            boot_split = boot_time.split()[3][:-1]
            uptime = int(round(time())) - int(boot_split)
        else:
            with open('/proc/uptime', 'r') as utime:
                uptime = utime.read().split('.')[0]

        self.key, self.value = self.__class__.__name__, f'{timedelta(seconds=int(uptime))}'[:-3]


class WindowManager(Utils):
    """Get the name of the Window Manager"""
    def __init__(self) -> None:
        wm_value, procs = 'Quartz' if self.is_platform('Darwin') else '', self.x_process()
        wm_value = [WM_DICT[k] for k in WM_DICT if k in procs][0]

        self.key, self.value = sub('(.*)([A-Z]+)', '\\1 \\2', self.__class__.__name__), wm_value


class DesktopEnvironment(Utils):
    """Get Desktop Environment"""
    def __init__(self) -> None:
        de_value, procs = '', self.x_process()
        de_value = [DE_DICT[k] for k in DE_DICT if k in procs][0]

        self.key, self.value = sub('(.*)([A-Z]+)', '\\1 \\2', self.__class__.__name__), de_value


class Shell(Utils):
    """Shell name"""
    def __init__(self) -> None:
        self.key, self.value = self.__class__.__name__, getenv('SHELL')


class Terminal(Utils):
    """Terminal name"""
    def __init__(self) -> None:
        self.key, self.value = self.__class__.__name__, getenv('TERM')


class Packages(Utils):
    """Get the number of installed packages"""
    def __init__(self) -> None:
        packages = ''

        # Remove rpm if dnf is installed, to avoid double runs
        PM_DICT.pop('rpm') if 'dnf' in PM_DICT else ''

        for pman in PM_DICT:
            if fExec(pman):
                PM_DICT[pman].insert(0, pman)
                pkgs0 = packages + '\n' + self.x_popen(PM_DICT[pman])
                pkgs1 = self.x_popen(PM_DICT[pman]).rstrip('\n')
                packages = pkgs0 if packages != '' else pkgs1

        packages = 'n/a' if packages == '' else len(packages.rstrip('\n').split('\n'))

        self.key, self.value = self.__class__.__name__, packages


class CPU(Utils):
    """Display the CPU information"""
    def __init__(self) -> None:
        cpu_info = 'n/a'
        if self.is_platform('Linux'):
            with open('/proc/cpuinfo') as cpu:
                cpu_info = sub('  +', ' ', cpu.readlines()[4])
        elif self.is_platform('Darwin'):
            cpu = self.x_popen(['sysctl', '-n', 'machdep.cpu.brand_string'])
            cpu_info = sub('  +', ' ', cpu)

        self.key, self.value = self.__class__.__name__, sub('model name\t: |CPU ', '', cpu_info).strip()


class RAM(Utils):
    """Calculate the RAM usage and format to display
    # RAM: usedMB, freeMB of totalGB"""
    def __init__(self) -> None:
        if self.is_platform('Linux'):
            raminfo = self.x_popen(['free', '-m', '--si']).split('\n')
            ram = raminfo[1].split()[1:]
            total_ram = f'{float(int(ram[0]) / 1000):.1f}'
            used_ram, free_ram = int(ram[1]), int(ram[2]) + int(ram[4])
        elif self.is_platform('Darwin'):
            mem_size = self.x_popen(['sysctl', '-n', 'hw.memsize'])
            total_ram = int(mem_size) / 1073741824

            get_memory = self.x_popen(['vm_stat']).replace('.', '')
            get_memory = sub(r':[\s]+', ' ', get_memory.replace('Pages ', '')).split('\n')
            mem = {}
            for row in range(1, len(get_memory) - 7):
                row = get_memory[row].replace('down', '').splitlines()[0].strip().split()
                mem[(row[0])] = int(row[1]) * 4000
            # @note: Using 4000 (eg 4096) makes it sum up exactly as in Activity Monitor

            used_ram = ceil((mem['wired'] + mem['active'] +
                             mem['inactive']) / 1048576 / 10.0) * 10
            free_ram = ceil((mem['free'] + mem['speculative']) / 1048576 / 10.0) * 10
        else:
            used_ram = free_ram = total_ram = '-'

        red, grn, clr = (COLOR_VARIES[0], COLOR_VARIES[1], COLOR_VARIES[-1])
        ram_display = f'{red}{used_ram}MB{clr}, {grn}{free_ram}MB{clr} of {total_ram}GB'

        self.key, self.value = self.__class__.__name__, ram_display


class Disk(Utils):
    """Calculate the disk usage and format to display
    # Disk: used (%) of total, N available (disk)"""
    def __init__(self) -> None:
        if self.is_platform('Linux'):
            filesystems = tuple(zip(('-t',) * len(FS_TYPES), FS_TYPES))
            fs_types = self.x_popen(list(chain(['df', '-HlPT', '--total'], *filesystems)))
            total = tuple(sub('%', '', t) for t in fs_types.splitlines()[-1].split())
            disknr, size, used, avail, capacity = total[1:-1]
            disknr = self.x_popen(['df', '/boot']).splitlines()[1].split()[0].replace('/dev/', '')
        elif self.is_platform('Darwin'):
            fs_types = self.x_popen(['df', '-HT', ','.join(FS_TYPES)])
            total = tuple(sub('(/dev/|%)', '', t) for t in fs_types.splitlines()[1].split())
            disknr, size, used, avail, capacity = total[:-1]

        used_percent = int(capacity) if capacity else '-'

        disk_color = 1
        disk_color = 2 if used_percent >= 50 else disk_color
        disk_color = 0 if used_percent >= 80 else disk_color
        if DISK_WARNING == 1 and (100 - used_percent) <= DISK_WARNING_THRESHOLD:
            disk_color = 3

        col, clr = (COLOR_VARIES[disk_color], COLOR_VARIES[-1])
        disk_display = f'{col}{used}B{clr} ({capacity}%) of {size}B, ' \
                       f'{avail}B available ({disknr})'

        self.key, self.value = self.__class__.__name__, disk_display


# Start output
OUT = Output()

LOGOS_DICT = {OUT.distro: '''
{color[1]}               +                {results[0]}
{color[1]}               #                {results[1]}
{color[1]}              ###               {results[2]}
{color[1]}             #####              {results[3]}
{color[1]}             ######             {results[4]}
{color[1]}            ; #####;            {results[5]}
{color[1]}           +##.#####            {results[6]}
{color[1]}          +##########           {results[7]}
{color[1]}         ######{color[0]}#####{color[1]}##;         {results[8]}
{color[1]}        ###{color[0]}############{color[1]}+        {results[9]}
{color[1]}       #{color[0]}######   #######        {results[10]}
{color[0]}     .######;     ;###;`\".      {results[11]}
{color[0]}    .#######;     ;#####.       {results[12]}
{color[0]}    #########.   .########`     {results[13]}
{color[0]}   ######'           '######    {results[14]}
{color[0]}  ;####                 ####;   {results[15]}
{color[0]}  ##'                     '##   {results[16]}
{color[0]} #'                         `#  {results[17]}
\x1b[0m'''}

# (OS X/macOS only) see --help
if OUT.distro == 'Darwin' and BRAND == 'Apple':
    LOGOS_DICT = {OUT.distro: '''
{color[1]}                     ##         {results[0]}
{color[1]}                  #####         {results[1]}
{color[1]}                 #####          {results[2]}
{color[1]}               #####            {results[3]}
{color[1]}               ###              {results[4]}
{color[1]}        #####      #####        {results[5]}
{color[1]}    #####################{color[0]}#      {results[6]}
{color[1]}  ###################{color[0]}#######    {results[7]}
{color[1]} ################{color[0]}#########      {results[8]}
{color[1]} #############{color[0]}###########       {results[9]}
{color[1]} ##########{color[0]}#############        {results[10]}
{color[1]} #######{color[0]}#################       {results[11]}
{color[1]} #####{color[0]}#####################     {results[12]}
{color[1]}  ###{color[0]}#######################    {results[13]}
{color[1]}   #{color[0]}#######################     {results[14]}
{color[0]}     #####################      {results[15]}
{color[0]}      ##################        {results[16]}
{color[0]}        ####      ####          {results[17]}
\x1b[0m'''}


# Display output
if __name__ == "__main__":
    for x in LABELS:
        OUT.append(vars()[CLASSES[x]]())
    print(OUT.output())
    sys.exit(0)
else:
    sys.exit(1)