Issue #5 feature commit

bunseokbot · Apr 14, 2019 · 3c81827 · 3c81827
1 parent 18ead57
commit 3c81827
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 12 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -3,7 +3,10 @@ MAINTAINER Namjun Kim <bunseokbot@gmail.com>
 
 ENV PYTHONUNBUFFERED 1
 
-RUN printf "deb http://archive.debian.org/debian/ jessie main\ndeb-src http://archive.debian.org/debian/ jessie main\ndeb http://security.debian.org jessie/updates main\ndeb-src http://security.debian.org jessie/updates main" > /etc/apt/sources.list
+RUN printf "deb http://archive.debian.org/debian/ jessie main"\
+           "\ndeb-src http://archive.debian.org/debian/ jessie main" \
+           "\ndeb http://security.debian.org jessie/updates main" \
+           "\ndeb-src http://security.debian.org jessie/updates main" > /etc/apt/sources.list
 
 RUN apt-get update \
     && apt-get install -y libmysqlclient-dev libfontconfig bzip2 curl \

diff --git a/config.ini b/config.ini
@@ -24,4 +24,4 @@ RESULT_BACKEND=redis://redis:6379/1
 AWS_ACCESS_KEY_ID=
 AWS_SECRET_ACCESS_KEY=
 REGION_NAME=
-BUCKET_NAME=darklight
+BUCKET_NAME=
diff --git a/crawler/__init__.py b/crawler/__init__.py
@@ -17,6 +17,7 @@
 import pipeline.source as pipelines
 
 import boto3
+import os
 
 
 class Crawler:
@@ -155,21 +156,31 @@ def save(self, id, obj):
             ).save()
 
     def upload_screenshot(self, screenshot, id):
-        """Upload screenshot into S3 storage."""
+        """Upload screenshot into S3 storage or local storage."""
         bucket = self.ini.read('STORAGE', 'BUCKET_NAME')
         key = f'screenshot/{id}.jpg'
 
-        client = boto3.client(service_name='s3',
-                              region_name=self.ini.read('STORAGE', 'REGION_NAME'),
-                              aws_access_key_id=self.ini.read('STORAGE', 'AWS_ACCESS_KEY_ID'),
-                              aws_secret_access_key=self.ini.read('STORAGE', 'AWS_SECRET_ACCESS_KEY'))
+        # if user want to upload screenshot into s3 storage
+        if bucket:
+            client = boto3.client(service_name='s3',
+                                  region_name=self.ini.read('STORAGE', 'REGION_NAME'),
+                                  aws_access_key_id=self.ini.read('STORAGE', 'AWS_ACCESS_KEY_ID'),
+                                  aws_secret_access_key=self.ini.read('STORAGE', 'AWS_SECRET_ACCESS_KEY'))
 
-        client.upload_fileobj(BytesIO(screenshot),
-                              Bucket=bucket,
-                              Key=key,
-                              ExtraArgs={'ACL': 'public-read'})
+            client.upload_fileobj(BytesIO(screenshot),
+                                  Bucket=bucket,
+                                  Key=key,
+                                  ExtraArgs={'ACL': 'public-read'})
 
-        return f"{client.meta.endpoint_url}/{bucket}/{key}"
+            return f"{client.meta.endpoint_url}/{bucket}/{key}"
+        else:
+            if not os.path.exists('screenshot'):
+                os.mkdir('screenshot')
+
+            with open(key, 'wb') as f:
+                f.write(screenshot)
+
+            return key
 
     def __del__(self):
         Log.i("Ending crawler")