Merge pull request #6 from becksteadn/build

use bucket object to upload to s3

src/lambda_function.py

@@ -42,16 +42,15 @@ def lambda_handler(event, context):
         url = 'http://' + url
     check_connection(url)
 
-    STAR_LEN = 40
     BUCKET_NAME = os.environ.get('GLIMPSE_BUCKET_NAME')
     SCREENSHOT_DIR = os.environ.get('GLIMPSE_SCREENSHOT_DIR')
     DB_TABLE = os.environ.get('GLIMPSE_DB_TABLE')
     pp = pprint.PrettyPrinter(indent=4)
 
     print('[!] Getting Environment Variables')
     print(f'Using S3 bucket "{BUCKET_NAME}"')
-    print(f'    with path "{SCREENSHOT_DIR}""')
-    print(f'Using DynamoDB table "{DB_TABLE}""')
+    print(f'    with path "{SCREENSHOT_DIR}"')
+    print(f'Using DynamoDB table "{DB_TABLE}"')
 
     now = datetime.now()
     timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
@@ -76,7 +75,6 @@ def lambda_handler(event, context):
 
     # Don't update if update==false or the parameter doesn't exist
     if 'update' not in event.keys() or str(event['update']).lower() != 'true':
-        # Don't force an update
         if exists:
             print('[!] Existing Data')
             pp.pprint(return_data)
@@ -85,11 +83,11 @@ def lambda_handler(event, context):
 
     glimpse = gd.GlimpseDriver()
     try:
+        print('[!] Rendering Page')
         glimpse.driver.get(url)
         glimpse.screenshot(local_path)
 
         s3 = S3(BUCKET_NAME)
-        #s3_key = s3.get_key(remote_path)
         s3.upload_file(local_path, remote_path)
 
         db_data['effectiveurl'] = glimpse.driver.current_url

src/s3_help.py

@@ -5,9 +5,12 @@ class S3:
     def __init__(self, bucket):
         self.bucket_name = bucket
         self.resource = boto3.resource('s3')
+        self.bucket = self.resource.Bucket(bucket)
 
     def get_key(self, remote_path):
         return self.resource.Object(bucket_name=self.bucket_name, key=remote_path)
 
     def upload_file(self, local_path, remote_path):
-        self.resource.meta.client.upload_file(local_path,self.bucket_name, remote_path)
+        print(f'[!] Uploading {local_path} to s3://{self.bucket_name}/{remote_path}')
+        return self.bucket.upload_file(local_path, remote_path)
+        #return self.resource.meta.client.upload_file(local_path,self.bucket_name, remote_path)