fix download_limit

app/routes.py

@@ -12,22 +12,25 @@ from apscheduler.schedulers.background import BackgroundScheduler
 import atexit
 
 
+def remove_share(sh):
+    # delete files
+    try:
+        rmtree(os.path.join(current_app.config['MEDIA_LOCATION'], sh.id))
+        os.remove(os.path.join(current_app.config['MEDIA_LOCATION'], sh.id + '.zip'))
+    except OSError as e:
+        print("Error: %s - %s." % (e.filename, e.strerror))
+    # delete from database
+    for att in sh.files:
+        db.session.delete(att)
+    db.session.delete(sh)
+    db.session.commit()
+
+
 def remove_expired():
-    print('CLEAING UP')
     with app.app_context():
         expired_shares = Share.query.filter(datetime.now() > Share.expires).all()
         for sh in expired_shares:
-            # delete files
-            try:
-                rmtree(os.path.join(current_app.config['MEDIA_LOCATION'], sh.id))
-                os.remove(os.path.join(current_app.config['MEDIA_LOCATION'], sh.id + '.zip'))
-            except OSError as e:
-                print("Error: %s - %s." % (e.filename, e.strerror))
-            # delete from database
-            for att in sh.files:
-                db.session.delete(att)
-            db.session.delete(sh)
-            db.session.commit()
+            remove_share(sh)
 
 
 sched = BackgroundScheduler(daemon=True)
@@ -36,24 +39,6 @@ sched.start()
 atexit.register(lambda: sched.shutdown())
 
 
-def check_limit(share):
-    if share.download_limit < 0 or share.expires < datetime.now():
-        # delete files
-        try:
-            rmtree(os.path.join(current_app.config['MEDIA_LOCATION'], share.id))
-            os.remove(os.path.join(current_app.config['MEDIA_LOCATION'], share.id + '.zip'))
-        except OSError as e:
-            print("Error: %s - %s." % (e.filename, e.strerror))
-        # delete from database
-        for att in share.files:
-            db.session.delete(att)
-        db.session.delete(share)
-        db.session.commit()
-        return False
-    else:
-        return True
-
-
 @app.route('/privacy')
 def privacy():
     return render_template('Privacy.html')
@@ -120,29 +105,38 @@ def upload():
 @app.route('/shared/<string:share_id>')
 def shared(share_id):
     share = Share.query.get_or_404(share_id)
-    check_limit(share)
     return render_template('Shared.html', url=url_for('download', share_id=share.id))
 
 
 @app.route('/d/<string:share_id>', methods=['GET'])
 def download(share_id):
     share = Share.query.get_or_404(share_id)
-    share.download_limit -= 1
-    db.session.commit()
-    return render_template('Download.html', up=share)
+    if share.download_limit < 0:
+        remove_share(share)
+        return redirect(url_for('expired'))
+    else:
+        return render_template('Download.html', up=share)
 
 
 @app.route('/media/<string:share_id>/<string:filename>', methods=['GET'])
 def media(share_id, filename):
     share = Share.query.get_or_404(share_id)
-    check_limit(share)
-    return send_from_directory(os.path.join(current_app.config['MEDIA_LOCATION'], share_id), filename=filename,
-                               as_attachment=True)
+    share.download_limit -= 1
+    db.session.commit()
+    if share.download_limit < 0:
+        remove_share(share)
+        return redirect(url_for('expired'))
+    else:
+        return send_from_directory(os.path.join(current_app.config['MEDIA_LOCATION'], share_id), filename=filename,
+                                   as_attachment=True)
 
 
 @app.route('/zip/<string:share_id>')
 def zip(share_id):
     share = Share.query.get_or_404(share_id)
-    check_limit(share)
-    return send_from_directory(os.path.join(current_app.config['MEDIA_LOCATION']),
-                               filename=str(share.id) + '.zip', as_attachment=True)
+    if share.download_limit < 0:
+        remove_share(share)
+        return redirect(url_for('expired'))
+    else:
+        return send_from_directory(os.path.join(current_app.config['MEDIA_LOCATION']),
+                                   filename=str(share.id) + '.zip', as_attachment=True)