diff --git a/snappass/main.py b/snappass/main.py index 8877a1e..00c4d01 100644 --- a/snappass/main.py +++ b/snappass/main.py @@ -5,7 +5,7 @@ import uuid import redis from cryptography.fernet import Fernet -from flask import abort, Flask, render_template, request +from flask import abort, Flask, render_template, request, jsonify from redis.exceptions import ConnectionError from werkzeug.urls import url_quote_plus from werkzeug.urls import url_unquote_plus @@ -177,7 +177,10 @@ def handle_password(): if URL_PREFIX: base_url = base_url + URL_PREFIX.strip("/") + "/" link = base_url + url_quote_plus(token) - return render_template('confirm.html', password_link=link) + if request.accept_mimetypes.accept_json and not request.accept_mimetypes.accept_html: + return jsonify(link=link, ttl=ttl) + else: + return render_template('confirm.html', password_link=link) @app.route('/', methods=['GET']) diff --git a/tests.py b/tests.py index 59459c0..1836801 100644 --- a/tests.py +++ b/tests.py @@ -3,6 +3,7 @@ import re import time import unittest import uuid +import json from unittest import TestCase from cryptography.fernet import Fernet @@ -138,6 +139,21 @@ class SnapPassRoutesTestCase(TestCase): frozen_time.move_to("2020-05-22 12:00:00") self.assertIsNone(snappass.get_password(key)) + + def test_set_password_json(self): + with freeze_time("2020-05-08 12:00:00") as frozen_time: + password = 'my name is my passport. verify me.' + rv = self.app.post('/', headers={'Accept': 'application/json'}, data={'password': password, 'ttl': 'two weeks'}) + + json_content = rv.get_json() + key = re.search(r'https://localhost/([^"]+)', json_content['link']).group(1) + key = unquote(key) + + frozen_time.move_to("2020-05-22 11:59:59") + self.assertEqual(snappass.get_password(key), password) + + frozen_time.move_to("2020-05-22 12:00:00") + self.assertIsNone(snappass.get_password(key)) if __name__ == '__main__':