gsutil

#!/usr/bin/env python3
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrap gsutil and perform local encryption and decryption.

Uses symmetric  keys stored in Cloud KMS.
"""

import random
import shutil
import string
import sys
import os

from encryption_wrapper import encryption
from encryption_wrapper.common import error_and_exit, run_command


# determine actual location of gsutil
_GSUTIL = os.getenv('GSUTIL_ACTUAL', '/snap/bin/gsutil')
# adding a random string to the tmp location for thread safety
random_str = ''.join(
    (random.choice(string.ascii_letters + string.digits) for i in range(8)))
_TMP_LOCATION = os.getenv(
    'GSUTIL_TMP_LOCATION',
    os.path.expanduser('~') + '/.gsutil-wrapper/' + random_str + '/')

class GSUtilWrapper(object):
  """Wrap the gsutil command to encrypt or decrypt files locally."""

  def __init__(self, argv):
    """Init for GSUtilWrapper

    Args:
      argv: gsutil command line arguments
    """
    self.argv = argv

  def wrap(self):
    """Wrap the gsutil command."""

    if self.argv[1] == 'cp' and \
        [i for i in self.argv if '--client_side_encryption' in i]:
      # if this is a cp command and we have the client side encryption argument
      # then proceed
      if 'linux' not in sys.platform:
        # not on a supported os
        error_and_exit(
            'You are running a wrapper around gsutil designed to handle local encryption/decryption transparently. Standard/original gsutil is available at {}'
            .format(_GSUTIL))

    elif self.argv[1] == 'cp' and self.argv[2] == '-r':
      error_and_exit(
          'encryption_wrapper wrapper does not yet support recursive copies. Please invoke {} directly.'
          .format(_GSUTIL))
    else:
      # run command without modification, exit with sig 0
      run_command('{} {}'.format(_GSUTIL, ' '.join(self.argv[1:])),
                  'gsutil command')
      sys.exit(0)

    # make a copy of gsutil arguments
    wrapped_args = self.argv.copy()
    to_url = wrapped_args[-1]
    from_url = wrapped_args[-2]

    # validation; can't locally encrypt if we're moving blobs between buckets
    if 'gs://' in to_url and 'gs://' in from_url:
      error_and_exit('cannot locally encrypt when from and two paths are in ' +
                     'the cloud')
    elif '*' in to_url or '*' in from_url:
      error_and_exit('wildcards are not yet supported')

    # grab our key_uri and creds strings from the arguments
    for arg in wrapped_args:
      if '--client_side_encryption' in arg:
        key_uri, creds = arg.split('=')[1].split(',')

    # the linter is worried these variables might not be initialized, but we
    # won't ever get this far if --client_side_encryption isn't specified
    # noinspection PyUnboundLocalVariable
    t = encryption.EncryptWithTink(key_uri, creds, _TMP_LOCATION)
    if 'gs://' in to_url:
      wrapped_args[-2] = t.encrypt(from_url)

    # now remove the client side encryption argument
    for i, arg in enumerate(wrapped_args):
      if '--client_side_encryption' in arg:
        del wrapped_args[i]

    # once the encryption/decryption is done, execute the gsutil command
    run_command(_GSUTIL + ' ' + ' '.join(wrapped_args[1:]),
                'wrapped gsutil command')
    if 'gs://' in from_url:
      if os.path.isdir(to_url):
        # need to append filename to target urls that are directories
        to_url = to_url + '/' + os.path.basename(from_url)
      t.decrypt(to_url)

    # set custom metadata
    if 'gs://' in to_url:
      if to_url.endswith('/'):
        # if the to url ends with a slash we need to grab the filename from
        # the from_url and append it
        object_url = to_url + from_url.split('/')[-1]
      else:
        object_url = to_url
      run_command(_GSUTIL +
                  ' setmeta -h "x-goog-meta-client-side-encrypted:true" ' +
                  object_url,
                  'set custom metadata')

    # clean up and exit with sig 0
    shutil.rmtree(_TMP_LOCATION)
    sys.exit(0)

def main():
  # we print this message so it's clear the user is talking to the wrapped
  # command and not gsutil itself
  print('gsutil is being wrapped. Standard gsutil available at: ' + _GSUTIL)

  try:
    wrapper = GSUtilWrapper(sys.argv)
    wrapper.wrap()
  except Exception as e:  # pylint disable=broad-except
    error_and_exit(str(e))

if __name__ == '__main__':
  main()