<commit_before>r"""Decompress scan files automatically on create.

Automatically decompress files in the gs://firehook-censoredplanetscanspublic/
bucket into the gs://firehook-scans/ bucket.

Files in the compressed bucket are of the form
gs://firehook-censoredplanetscanspublic/CP_Quack-discard-2018-07-28-03-11-21.tar.gz
or
gs://firehook-censoredplanetscanspublic/CP_Satellite-2018-08-07-17-24-41.tar.gz

They should be decompressed into scan type specific directories like
gs://firehook-scans/discard/CP_Quack-discard-2018-07-28-03-11-21/results.json
or
gs://firehook-scans/satellite/CP_Satellite-2018-08-07-17-24-41.tar.gz


To deploy this cloud function run

gcloud functions deploy decompress_scans \
  --runtime python37 \
  --region us-east1 \
  --project firehook-censoredplanet \
  --trigger-resource firehook-censoredplanetscanspublic \
  --trigger-event google.storage.object.finalize

View deployed functions at
https://pantheon.corp.google.com/functions/list?project=firehook-censoredplanet
"""

import io
import os
import tarfile
from pprint import pprint

from google.cloud import storage

COMPRESSED_BUCKET_NAME = 'firehook-censoredplanetscanspublic'
UNCOMPRESSED_BUCKET_NAME = 'firehook-scans'

client = storage.Client()
compressed_bucket = client.get_bucket(COMPRESSED_BUCKET_NAME)
uncompressed_bucket = client.get_bucket(UNCOMPRESSED_BUCKET_NAME)

# Key - a substring that will be in the filename
# Value - the scan type of that file
scan_type_identifiers = {
    'Satellite-': 'satellite',
    'Quack-echo-': 'echo',
    'Quack-discard-': 'discard',
    'Quack-https-': 'https',
    'Quack-http-': 'http'
}


def decompress_scans(event, context):
  """Decompress files on cloud storage creates.

  Called by the cloud function.

  Args:
    event: a dictionary of info on the triggering event
      format https://cloud.google.com/storage/docs/json_api/v1/objects#resource
    context: metadata on the triggering event
  """
  tar_name = event['name']
  decompress_file(tar_name)


def decompress_file(tar_name):
  """Decompress a given scan file.

  Args:
    tar_name: filename like CP_Quack-discard-2020-08-17-08-41-15.tar.gz

  Raises:
    Exception: when the file has an unknown scan type
  """
  scan_type = None
  for type_identifier, potential_scan_type in scan_type_identifiers.items():
    if type_identifier in tar_name:
      scan_type = potential_scan_type

  if not scan_type:
    raise Exception("Couldn't determine scan type for filename " + tar_name)

  input_blob = compressed_bucket.get_blob(tar_name).download_as_string()
  tar = tarfile.open(fileobj=io.BytesIO(input_blob))

  for file_path in tar.getnames():
    file_object = tar.extractfile(file_path)

    # skip directories
    if file_object:
      output_blob = uncompressed_bucket.blob(os.path.join(scan_type, file_path))
      output_blob.upload_from_string(file_object.read())


def get_all_compressed_filenames():
  """Get a list of all compressed filenames, minus the file extension.

  Returns:
    a list of filename strings
    ex ["CP_Quack-discard-2020-08-17-08-41-15",
        "CP_Satellite-2020-08-16-17-07-54"]
  """
  # CP_Satellite-2020-08-16-17-07-54.tar.gz
  blobs = list(client.list_blobs(COMPRESSED_BUCKET_NAME))
  # CP_Satellite-2020-08-16-17-07-54
  filenames = [  # remove both .tar and .gz
      os.path.splitext(os.path.splitext(blob.name)[0])[0] for blob in blobs
  ]
  return filenames


def get_all_uncompressed_filepaths():
  """Get a list of all directories with uncompressed filenames.

  Returns:
    a list of filename strings
    ex ["CP_Quack-discard-2020-08-17-08-41-15",
        "CP_Satellite-2020-08-16-17-07-54"]
  """
  # discard/CP_Quack-discard-2020-08-17-08-41-15/results.json
  blobs = list(client.list_blobs(UNCOMPRESSED_BUCKET_NAME))
  # discard/CP_Quack-discard-2020-08-17-08-41-15/
  paths = [os.path.split(blob.name)[0] for blob in blobs]
  # CP_Quack-discard-2020-08-17-08-41-15/
  path_ends = [os.path.split(path)[1] for path in paths]
  return path_ends


def get_missing_compressed_files(compressed_files, uncompressed_files):
  """Get all files in the compressed list that are not in the uncompressed list."""
  diff = set(compressed_files) - set(uncompressed_files)
  return list(diff)


def decompress_all_missing_files():
  """Decompress all files that exist only in the compressed bucket.

  Used for backfilling data.
  """

  compressed_files = get_all_compressed_filenames()
  uncompressed_files = get_all_uncompressed_filepaths()
  new_files = get_missing_compressed_files(compressed_files, uncompressed_files)

  files_with_extensions = [filename + '.tar.gz' for filename in new_files]

  for filename in files_with_extensions:
    #pprint(filename)
    decompress_file(filename)


if __name__ == '__main__':
  # Called manually
  decompress_all_missing_files()<commit_msg>cleanup<commit_after>r"""Decompress scan files automatically on create.

Automatically decompress files in the gs://firehook-censoredplanetscanspublic/
bucket into the gs://firehook-scans/ bucket.

Files in the compressed bucket are of the form
gs://firehook-censoredplanetscanspublic/CP_Quack-discard-2018-07-28-03-11-21.tar.gz
or
gs://firehook-censoredplanetscanspublic/CP_Satellite-2018-08-07-17-24-41.tar.gz

They should be decompressed into scan type specific directories like
gs://firehook-scans/discard/CP_Quack-discard-2018-07-28-03-11-21/results.json
or
gs://firehook-scans/satellite/CP_Satellite-2018-08-07-17-24-41.tar.gz


To deploy this cloud function run

gcloud functions deploy decompress_scans \
  --runtime python37 \
  --region us-east1 \
  --project firehook-censoredplanet \
  --trigger-resource firehook-censoredplanetscanspublic \
  --trigger-event google.storage.object.finalize

View deployed functions at
https://pantheon.corp.google.com/functions/list?project=firehook-censoredplanet
"""

import io
import os
import tarfile

from google.cloud import storage

COMPRESSED_BUCKET_NAME = 'firehook-censoredplanetscanspublic'
UNCOMPRESSED_BUCKET_NAME = 'firehook-scans'

client = storage.Client()
compressed_bucket = client.get_bucket(COMPRESSED_BUCKET_NAME)
uncompressed_bucket = client.get_bucket(UNCOMPRESSED_BUCKET_NAME)

# Key - a substring that will be in the filename
# Value - the scan type of that file
scan_type_identifiers = {
    'Satellite-': 'satellite',
    'Quack-echo-': 'echo',
    'Quack-discard-': 'discard',
    'Quack-https-': 'https',
    'Quack-http-': 'http'
}


def decompress_scans(event, context):
  """Decompress files on cloud storage creates.

  Called by the cloud function.

  Args:
    event: a dictionary of info on the triggering event
      format https://cloud.google.com/storage/docs/json_api/v1/objects#resource
    context: metadata on the triggering event
  """
  tar_name = event['name']
  decompress_file(tar_name)


def decompress_file(tar_name):
  """Decompress a given scan file.

  Args:
    tar_name: filename like CP_Quack-discard-2020-08-17-08-41-15.tar.gz

  Raises:
    Exception: when the file has an unknown scan type
  """
  scan_type = None
  for type_identifier, potential_scan_type in scan_type_identifiers.items():
    if type_identifier in tar_name:
      scan_type = potential_scan_type

  if not scan_type:
    raise Exception("Couldn't determine scan type for filename " + tar_name)

  input_blob = compressed_bucket.get_blob(tar_name).download_as_string()
  tar = tarfile.open(fileobj=io.BytesIO(input_blob))

  for file_path in tar.getnames():
    file_object = tar.extractfile(file_path)

    # skip directories
    if file_object:
      output_blob = uncompressed_bucket.blob(os.path.join(scan_type, file_path))
      output_blob.upload_from_string(file_object.read())


def get_all_compressed_filenames():
  """Get a list of all compressed filenames, minus the file extension.

  Returns:
    a list of filename strings
    ex ["CP_Quack-discard-2020-08-17-08-41-15",
        "CP_Satellite-2020-08-16-17-07-54"]
  """
  # CP_Satellite-2020-08-16-17-07-54.tar.gz
  blobs = list(client.list_blobs(COMPRESSED_BUCKET_NAME))
  # CP_Satellite-2020-08-16-17-07-54
  filenames = [  # remove both .tar and .gz
      os.path.splitext(os.path.splitext(blob.name)[0])[0] for blob in blobs
  ]
  return filenames


def get_all_uncompressed_filepaths():
  """Get a list of all directories with uncompressed filenames.

  Returns:
    a list of filename strings
    ex ["CP_Quack-discard-2020-08-17-08-41-15",
        "CP_Satellite-2020-08-16-17-07-54"]
  """
  # discard/CP_Quack-discard-2020-08-17-08-41-15/results.json
  blobs = list(client.list_blobs(UNCOMPRESSED_BUCKET_NAME))
  # discard/CP_Quack-discard-2020-08-17-08-41-15/
  paths = [os.path.split(blob.name)[0] for blob in blobs]
  # CP_Quack-discard-2020-08-17-08-41-15/
  path_ends = [os.path.split(path)[1] for path in paths]
  return path_ends


def get_missing_compressed_files(compressed_files, uncompressed_files):
  """Get all files in the compressed list that are not in the uncompressed list."""
  diff = set(compressed_files) - set(uncompressed_files)
  return list(diff)


def decompress_all_missing_files():
  """Decompress all files that exist only in the compressed bucket.

  Used for backfilling data.
  """
  compressed_files = get_all_compressed_filenames()
  uncompressed_files = get_all_uncompressed_filepaths()
  new_files = get_missing_compressed_files(compressed_files, uncompressed_files)

  files_with_extensions = [filename + '.tar.gz' for filename in new_files]

  for filename in files_with_extensions:
    decompress_file(filename)


if __name__ == '__main__':
  # Called manually
  decompress_all_missing_files()
{'fraction_non_alphanumeric': 0.10558618537397274, 'fraction_numerical': 0.0424425257463851, 'mean_word_length': 4.262178434592228, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1988457', 'n_tokens_mistral': 3626, 'n_tokens_neox': 3274, 'n_words': 712}