diff --git a/app.py b/app.py index 7a74f2b..8614794 100644 --- a/app.py +++ b/app.py @@ -4,6 +4,7 @@ A simple Flask application for interacting with WHSPAH shockers. import argparse import logging +import random import json import os @@ -58,19 +59,22 @@ def transmit(): txID = int(data['transmitterID']) channel = int(data['channel']) action = ACTION_MAP[data['action']] + pin = int(data['shockerPin']) intensity = int(data.get('intensity', 0)) lucal = bool(data.get('lucalEncoded', False)) - pin = int(data['shockerPin']) # If any of those failed, return an error. except (ValueError, KeyError): return {'success': False, 'message': 'Request must contain the following keys:\n'+ 'txID: int,\n'+ 'channel: int,\n'+ - 'pin: int,\n'+ 'action: "shock", "vibrate" or "beep",\n'+ + 'pin: int,\n'+ 'intensity (optional): int\n'+ 'lucalEncoded (optional): bool'}, 400 + if pin != app.config['pin']: + return {'success': False, 'message': 'Unauthorised'}, 401 + # Send the data to WHSPAH tx: whspah.Transmitter = app.config['transmitter'] tx.transmit(txID, channel, action, intensity, lucal) @@ -85,6 +89,9 @@ if __name__ == '__main__': parser.add_argument('-d', '--debug', action='store_true') parser.add_argument('--ip', type=str, default='0.0.0.0') parser.add_argument('--port', type=int, default=8000) + parser.add_argument('--pin', type=int, default=random.randint(1000, 9999), + help='The authorisation pin, defaults to a random number'+ + 'between 1000 and 9999') args = parser.parse_args() # Sets the logging version based on --debug @@ -96,6 +103,10 @@ if __name__ == '__main__': # Connect to a WHSPAH device app.config['transmitter'] = whspah.Transmitter() + # Configure pin + app.config['pin'] = args.pin + print(f'Running with pin {args.pin}.') + if args.debug: # If in debug mode, run Flask's built in server app.run(host=args.ip, port=args.port, debug=True)