7 files changed, 62 insertions, 13 deletions
diff --git a/app/tasks/importtasks.py b/app/tasks/importtasks.py
index e0102b8..638e9ca 100644
--- a/app/tasks/importtasks.py
+++ b/app/tasks/importtasks.py
@@ -1,10 +1,12 @@
-import flask, json
+import flask, json, os
 from flask.ext.sqlalchemy import SQLAlchemy
+from urllib.error import HTTPError
 import urllib.request
 from urllib.parse import urlparse, quote_plus
 from app import app
 from app.models import *
 from app.tasks import celery, TaskError
+from app.utils import randomString
 
 class GithubURLMaker:
 	def __init__(self, url):
@@ -37,7 +39,7 @@ class GithubURLMaker:
 		return self.baseUrl + "/description.txt"
 
 	def getScreenshotURL(self):
-		return self.baseUrl + "/placeholder.png"
+		return self.baseUrl + "/screenshot.png"
 
 	def getCommitsURL(self, branch):
 		return "https://api.github.com/repos/{}/{}/commits?sha={}" \
@@ -147,7 +149,7 @@ def getMeta(urlstr, author):
 				result[key] = conf[key]
 			except KeyError:
 				pass
-	except OSError:
+	except HTTPError:
 		print("mod.conf does not exist")
 
 	if "name" in result:
@@ -157,7 +159,7 @@ def getMeta(urlstr, author):
 		try:
 			contents = urllib.request.urlopen(urlmaker.getDescURL()).read().decode("utf-8")
 			result["description"] = contents.strip()
-		except OSError:
+		except HTTPError:
 			print("description.txt does not exist!")
 
 	if "description" in result:
@@ -172,6 +174,7 @@ def getMeta(urlstr, author):
 
 	return result
 
+
 @celery.task()
 def makeVCSRelease(id, branch):
 	release = PackageRelease.query.get(id)
@@ -204,3 +207,40 @@ def makeVCSRelease(id, branch):
 	db.session.commit()
 
 	return release.url
+
+
+@celery.task()
+def importRepoScreenshot(id):
+	package = Package.query.get(id)
+	if package is None:
+		raise Exception("Unexpected none package")
+
+	# Get URL Maker
+	url = urlparse(package.repo)
+	urlmaker = None
+	if url.netloc == "github.com":
+		urlmaker = GithubURLMaker(url)
+	else:
+		raise TaskError("Unsupported repo")
+
+	if not urlmaker.isValid():
+		raise TaskError("Error! Url maker not valid")
+
+	try:
+		filename = randomString(10) + ".png"
+		imagePath = os.path.join(app.config["UPLOAD_FOLDER"], filename)
+		print(imagePath)
+		urllib.request.urlretrieve(urlmaker.getScreenshotURL(), imagePath)
+
+		ss = PackageScreenshot()
+		ss.package = package
+		ss.title   = "screenshot.png"
+		ss.url     = "/uploads/" + filename
+		db.session.add(ss)
+		db.session.commit()
+
+		return "/uploads/" + filename
+	except HTTPError:
+		print("screenshot.png does not exist")
+
+	return None
diff --git a/app/views/utils.py b/app/utils.py
index 44647a1..44647a1 100644
--- a/app/views/utils.py
+++ b/app/utils.py
diff --git a/app/views/admin.py b/app/views/admin.py
index ee4d9f6..03dcc75 100644
--- a/app/views/admin.py
+++ b/app/views/admin.py
@@ -6,7 +6,7 @@ from app.models import *
 from app.tasks.forumtasks import importUsersFromModList
 from flask_wtf import FlaskForm
 from wtforms import *
-from .utils import loginUser, rank_required
+from app.utils import loginUser, rank_required
 
 @menu.register_menu(app, ".admin", "Admin", order=30,
 		visible_when=lambda: current_user.rank.atLeast(UserRank.ADMIN))
diff --git a/app/views/githublogin.py b/app/views/githublogin.py
index 114bb2b..27b197c 100644
--- a/app/views/githublogin.py
+++ b/app/views/githublogin.py
@@ -5,7 +5,7 @@ import flask_menu as menu
 from flask_github import GitHub
 from app import app, github
 from app.models import *
-from .utils import loginUser
+from app.utils import loginUser
 
 @app.route("/user/github/start/")
 def github_signin_page():
diff --git a/app/views/packages.py b/app/views/packages.py
index 4a5fd3b..f2b17f1 100644
--- a/app/views/packages.py
+++ b/app/views/packages.py
@@ -3,10 +3,11 @@ from flask_user import *
 from flask.ext import menu
 from app import app
 from app.models import *
-from app.tasks.importtasks import makeVCSRelease
+from app.tasks.importtasks import importRepoScreenshot, makeVCSRelease
 
-from .utils import *
+from app.utils import *
 
+from urllib.parse import urlparse
 from flask_wtf import FlaskForm
 from wtforms import *
 from wtforms.validators import *
@@ -153,10 +154,11 @@ def create_edit_package_page(author=None, name=None):
 
 	# Initial form class from post data and default data
 	if request.method == "POST" and form.validate():
-		# Successfully submitted!
+		wasNew = False
 		if not package:
 			package = Package()
 			package.author = author
+			wasNew = True
 		else:
 			triggerNotif(package.author, current_user,
 					"{} edited".format(package.title), package.getDetailsURL())
@@ -168,7 +170,14 @@ def create_edit_package_page(author=None, name=None):
 			package.tags.append(Tag.query.get(tag))
 
 		db.session.commit() # save
-		return redirect(package.getDetailsURL()) # redirect
+
+		if wasNew:
+			url = urlparse(package.repo)
+			if url.netloc == "github.com":
+				task = importRepoScreenshot.delay(package.id)
+				return redirect(url_for("check_task", id=task.id, r=package.getDetailsURL()))
+
+		return redirect(package.getDetailsURL())
 
 	return render_template("packages/create_edit.html", package=package, form=form, author=author)
 
diff --git a/app/views/tasks.py b/app/views/tasks.py
index 3759543..150fb81 100644
--- a/app/views/tasks.py
+++ b/app/views/tasks.py
@@ -5,10 +5,10 @@ from app import app, csrf
 from app.models import *
 from app.tasks import celery, TaskError
 from app.tasks.importtasks import getMeta
-from .utils import shouldReturnJson
+from app.utils import shouldReturnJson
 # from celery.result import AsyncResult
 
-from .utils import *
+from app.utils import *
 
 @csrf.exempt
 @app.route("/tasks/getmeta/new/", methods=["POST"])
diff --git a/app/views/users.py b/app/views/users.py
index ed195d4..19b0442 100644
--- a/app/views/users.py
+++ b/app/views/users.py
@@ -8,7 +8,7 @@ from flask_wtf import FlaskForm
 from flask_user.forms import RegisterForm
 from wtforms import *
 from wtforms.validators import *
-from .utils import rank_required, randomString
+from app.utils import rank_required, randomString
 from app.tasks.forumtasks import checkForumAccount
 from app.tasks.emails import sendVerifyEmail