WizardKit/scripts/wk/debug.py

"""WizardKit: Debug Functions"""
# vim: sts=2 sw=2 ts=2

import inspect
import logging
import lzma
import os
import pathlib
import pickle
import platform
import re
import socket
import sys
import time

from typing import Any

import requests

from wk.cfg.net import CRASH_SERVER
from wk.log import get_root_logger_path

# Classes
class Debug():
  """Object used when dumping debug data."""
  def method(self) -> None:
    """Dummy method used to identify functions vs data."""


# STATIC VARIABLES
LOG = logging.getLogger(__name__)
DEBUG_CLASS = Debug()
METHOD_TYPE = type(DEBUG_CLASS.method)


# Functions
def generate_debug_report() -> str:
  """Generate debug report, returns str."""
  platform_function_list = (
    'architecture',
    'machine',
    'platform',
    'python_version',
    )
  report = []

  # Logging data
  try:
    log_path = get_root_logger_path()
  except RuntimeError:
    # Assuming logging wasn't started
    pass
  else:
    report.append('------ Start Log -------')
    report.append('')
    with open(log_path, 'r', encoding='utf-8') as log_file:
      report.extend(log_file.read().splitlines())
    report.append('')
    report.append('------- End Log --------')

  # System
  report.append('--- Start debug info ---')
  report.append('')
  report.append('[System]')
  report.append(f'  {"FQDN":<24} {socket.getfqdn()}')
  for func in platform_function_list:
    func_name = func.replace('_', ' ').capitalize()
    func_result = getattr(platform, func)()
    report.append(f'  {func_name:<24} {func_result}')
  report.append(f'  {"Python sys.argv":<24} {sys.argv}')
  report.append('')

  # Environment
  report.append('[Environment Variables]')
  for key, value in sorted(os.environ.items()):
    report.append(f'  {key:<24} {value}')
  report.append('')

  # Done
  report.append('---- End debug info ----')
  return '\n'.join(report)


def generate_object_report(obj: Any, indent: int = 0) -> list[str]:
  """Generate debug report for obj, returns list."""
  report = []
  attr_list = []

  # Get attribute list
  if hasattr(obj, '__slots__'):
    attr_list = list(obj.__slots__)
  else:
    attr_list = [name for name in dir(obj) if not name.startswith('_')]

  # Dump object data
  for name in attr_list:
    attr = getattr(obj, name)

    # Skip methods
    if isinstance(attr, METHOD_TYPE):
      continue

    # Add attribute to report (expanded if necessary)
    if isinstance(attr, dict):
      report.append(f'{name}:')
      for key, value in attr.items():
        report.append(f'{"  "*(indent+1)}{key}: {str(value)}')
    else:
      report.append(f'{"  "*indent}{name}: {str(attr)}')

  # Done
  return report


def save_pickles(
    obj_dict: dict[Any, Any],
    out_path: pathlib.Path | str | None = None,
    ) -> None:
  """Save dict of objects using pickle."""
  LOG.info('Saving pickles')

  # Set path
  if not out_path:
    out_path = get_root_logger_path()
    out_path = out_path.parent.joinpath('../debug').resolve()

  # Save pickles
  try:
    for name, obj in obj_dict.copy().items():
      if name.startswith('__') or inspect.ismodule(obj):
        continue
      with open(f'{out_path}/{name}.pickle', 'wb') as _f:
        pickle.dump(obj, _f, protocol=pickle.HIGHEST_PROTOCOL)
  except Exception:
    LOG.error('Failed to save all the pickles', exc_info=True)


def upload_debug_report(
    report: str,
    compress: bool = True,
    reason: str = 'DEBUG',
    ) -> None:
  """Upload debug report to CRASH_SERVER as specified in wk.cfg.main."""
  LOG.info('Uploading debug report to %s', CRASH_SERVER.get('Name', '?'))
  headers = CRASH_SERVER.get('Headers', {'X-Requested-With': 'XMLHttpRequest'})
  if compress:
    headers['Content-Type'] = 'application/octet-stream'

  # Check if the required server details are available
  if not all(CRASH_SERVER.get(key, False) for key in ('Name', 'Url', 'User')):
    msg = 'Server details missing, aborting upload.'
    print(msg)
    raise RuntimeError(msg)

  # Set filename (based on the logging config if possible)
  filename = 'Unknown'
  try:
    log_path = get_root_logger_path()
  except RuntimeError:
    # Assuming logging wasn't started
    pass
  else:
    # Strip everything but the prefix
    filename = re.sub(r'^(.*)_(\d{4}-\d{2}-\d{2}.*)', r'\1', log_path.name)
  filename = f'{filename}_{reason}_{time.strftime("%Y-%m-%d_%H%M%S%z")}.log'
  LOG.debug('filename: %s', filename)

  # Compress report
  if compress:
    filename += '.xz'
    xz_report = lzma.compress(report.encode('utf8'))

  # Upload report
  url = f'{CRASH_SERVER["Url"]}/{filename}'
  response = requests.put(
    url,
    auth=(CRASH_SERVER['User'], CRASH_SERVER.get('Pass', '')),
    data=xz_report if compress else report,
    headers=headers,
    timeout=60,
    )

  # Check response
  if not response.ok:
    raise RuntimeError('Failed to upload report')


if __name__ == '__main__':
  print("This file is not meant to be called directly.")