runner_utils.py 4.45 KB
Newer Older
1 2 3 4
import errno
import glob
import os
import subprocess
5
import sys
6 7 8 9

from contextlib import contextmanager
from hashlib import sha256
from zc.buildout.configparser import parse
10 11 12
from slapos.util import bytes2str, str2bytes

import six
13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68


@contextmanager
def CwdContextManager(path):
  """
  Context Manager for executing code in a given directory.
  There is no need to provide fallback or basic checks
  in this code, as these checkes should exist in the code
  invoking this Context Manager.
  If someone needs to add checks here, I'm pretty sure
  it means that they are trying to hide legitimate errors.
  See tests to see examples of invokation
  """
  old_path = os.getcwd()
  try:
    os.chdir(path)
    yield
  finally:
    os.chdir(old_path)


def getExcludePathList(path):
  excluded_path_list = [
    "*.sock",
    "*.socket",
    "*.pid",
    ".installed*.cfg",
  ]

  def append_relative(path_list):
    for p in path_list:
      p = p.strip()
      if p:
        excluded_path_list.append(os.path.relpath(p, path))

  for partition in glob.glob(os.path.join(path, "instance", "slappart*")):
    if not os.path.isdir(partition):
      continue

    with CwdContextManager(partition):
      try:
        with open("srv/exporter.exclude") as f:
          exclude = f.readlines()
      except IOError as e:
        if e.errno != errno.ENOENT:
          raise
      else:
        append_relative(exclude)
      for installed in glob.glob(".installed*.cfg"):
        try:
          with open(installed) as f:
            installed = parse(f, installed)
        except IOError as e:
          if e.errno != errno.ENOENT:
            raise
        else:
69
          for section in six.itervalues(installed):
70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
            append_relative(section.get(
              '__buildout_installed__', '').splitlines())

  return excluded_path_list


def getSlappartSignatureMethodDict():
  slappart_signature_method_dict = {}
  for partition in glob.glob("./instance/slappart*"):
    if os.path.isdir(partition):
      script_path = os.path.join(partition, 'srv', '.backup_identity_script')
      if os.path.exists(script_path):
        slappart_signature_method_dict[partition] = script_path
  return slappart_signature_method_dict


def read_file_by_chunk(path, chunk_size=1024 * 1024):
  with open(path, 'rb') as f:
    chunk = f.read(chunk_size)
    while chunk:
      yield chunk
      chunk = f.read(chunk_size)


def getSha256Sum(file_path_list):
  result_list = []

  for file_path in file_path_list:
    hash_sum = sha256()
    for chunk in read_file_by_chunk(file_path):
      hash_sum.update(chunk)
    result_list.append("%s  %s" % (hash_sum.hexdigest(), file_path))

  return result_list


def writeSignatureFile(slappart_signature_method_dict, runner_working_path, signature_file_path='backup.signature'):
  special_slappart_list = slappart_signature_method_dict.keys()
  signature_list = []

  for dirpath, dirname_list, filename_list in os.walk('.'):
    if dirpath == '.' or not filename_list:
      continue

    # Find if special signature function should be applied
115
    signature_process = None
116 117 118 119 120 121
    for special_slappart in special_slappart_list:
      backup_identity_script_path = os.path.join(
        runner_working_path,
        slappart_signature_method_dict[special_slappart]
      )

122
      if dirpath.startswith('./' + os.path.relpath(os.path.join('./runner', special_slappart))):
123 124 125 126 127
        signature_process = subprocess.Popen(
          backup_identity_script_path,
          stdin=subprocess.PIPE,
          stdout=subprocess.PIPE,
        )
128
        break
129

130 131 132
    # construct list of file path and remove broken symlink
    filepath_list = filter(os.path.isfile, [os.path.join(dirpath, filename) for filename in filename_list])

133 134
    if signature_process:
      (output, error_output) = signature_process.communicate(
135
        str2bytes('\0'.join(filepath_list))
136
      )
137

138 139 140 141 142
      if signature_process.returncode != 0:
        print(
          "An issue occured when calculating the custom signature"
          " with %s :\n%s\n%s" % (
            backup_identity_script_path, output, error_output
143 144
          )
        )
145 146 147 148
        sys.exit(1)

      # We have to rstrip as most programs return an empty line
      # at the end of their output
149
      signature_list.extend(bytes2str(output).strip('\n').split('\n'))
150 151
    else:
      signature_list.extend(
152
        getSha256Sum(filepath_list)
153
      )
154

155 156 157
  # Write the signatures in file
  with open(signature_file_path, 'w+') as signature_file:
    signature_file.write("\n".join(sorted(signature_list)))