WizardKit/scripts/wk/debug.py

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

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

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):
    """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():
  """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, indent=0):
  """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, out_path=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, compress=True, reason='DEBUG'):
  """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.")