import boto3
import os
import json

session = boto3.Session(
  aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'],
  aws_secret_access_key=os.environ['AWS_SECRET_ACCESS_KEY']
)

s3 = session.resource('s3', region_name=os.environ['AWS_REGION'])
sqs = session.client('sqs', region_name=os.environ['AWS_REGION'])

def get_key_from_message(body):
  return body['Records'][0]['s3']['object']['key']

def get_bucket_from_message(body):
  return body['Records'][0]['s3']['bucket']['name']

def get_object(key, bucket):
  return s3.Object(
    bucket_name=bucket,
    key=key
  ).get()['Body'].read()

def file_stream_upload(buffer, name, bucket):
  return s3.Bucket(bucket).upload_fileobj(buffer, name)

def receive_message(queue, message_num=1, wait_time=1):
  return sqs.receive_message(
    QueueUrl=queue,
    MaxNumberOfMessages=message_num,
    WaitTimeSeconds=wait_time
  )

def delete_message(queue, receipt):
  return sqs.delete_message(
    QueueUrl=queue,
    ReceiptHandle=receipt
  )