diff --git a/aw_server/api.py b/aw_server/api.py index cc060932..2711be8f 100644 --- a/aw_server/api.py +++ b/aw_server/api.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Any, Optional +from typing import Dict, List, Any, Optional, Union from datetime import datetime from socket import gethostname import functools @@ -33,7 +33,7 @@ class ServerAPI: def __init__(self, db, testing) -> None: self.db = db self.testing = testing - self.last_event = {} #type: dict + self.last_event = {} # type: dict def get_info(self) -> Dict[str, Dict]: """Get server info""" @@ -73,11 +73,41 @@ def export_bucket(self, bucket_id: str) -> Dict[str, Any]: del event["id"] return bucket - def create_bucket(self, bucket_id: str, event_type: str, client: str, hostname: str) -> bool: + def export_all(self) -> Dict[str, Any]: + """Exports all buckets and their events to a format consistent across versions""" + buckets = self.get_buckets() + exported_buckets = {} + for bid in buckets.keys(): + exported_buckets[bid] = self.export_bucket(bid) + return exported_buckets + + def import_bucket(self, bucket_data: Any): + bucket_id = bucket_data["id"] + logger.info("Importing bucket {}".format(bucket_id)) + # TODO: Check that bucket doesn't already exist + self.db.create_bucket( + bucket_id, + type=bucket_data["type"], + client=bucket_data["client"], + hostname=bucket_data["hostname"], + created=(bucket_data["created"] + if isinstance(bucket_data["created"], datetime) + else iso8601.parse_date(bucket_data["created"])), + ) + self.create_events(bucket_id, [Event(**e) if isinstance(e, dict) else e for e in bucket_data["events"]]) + + def import_all(self, buckets: Dict[str, Any]): + for bid, bucket in buckets.items(): + self.import_bucket(bucket) + + def create_bucket(self, bucket_id: str, event_type: str, client: str, + hostname: str, created: Optional[datetime] = None)-> bool: """ Create bucket. Returns True if successful, otherwise false if a bucket with the given ID already existed. """ + if created is None: + created = datetime.now() if bucket_id in self.db.buckets(): return False self.db.create_bucket( @@ -85,7 +115,7 @@ def create_bucket(self, bucket_id: str, event_type: str, client: str, hostname: type=event_type, client=client, hostname=hostname, - created=datetime.now() + created=created ) return True @@ -207,10 +237,3 @@ def get_log(self): for line in log_file.readlines()[::-1]: payload.append(json.loads(line)) return payload, 200 - - def export_all(self) -> Dict[str, dict]: - """Exports all buckets and their events to a format consistent across versions""" - buckets = self.get_buckets() - for bid in buckets.keys(): - buckets[bid] = self.export_bucket(bid) - return buckets diff --git a/aw_server/rest.py b/aw_server/rest.py index ff925a73..f558c819 100644 --- a/aw_server/rest.py +++ b/aw_server/rest.py @@ -2,7 +2,7 @@ import traceback import json -from flask import request, Blueprint, jsonify +from flask import request, Blueprint, jsonify, make_response from flask_restplus import Api, Resource, fields import iso8601 from datetime import datetime, timedelta @@ -26,6 +26,7 @@ blueprint = Blueprint('api', __name__, url_prefix='/api') api = Api(blueprint, doc='/') + # TODO: Clean up JSONEncoder code? class CustomJSONEncoder(json.JSONEncoder): def __init__(self, *args, **kwargs): @@ -40,6 +41,8 @@ def default(self, obj, *args, **kwargs): except TypeError: pass return json.JSONEncoder.default(self, obj) + + app.json_encoder = CustomJSONEncoder app.register_blueprint(blueprint) @@ -54,26 +57,20 @@ def format(self, value): return json.loads(value) -# Loads event schema from JSONSchema in aw_core +# Loads event and bucket schema from JSONSchema in aw_core event = api.schema_model('Event', schema.get_json_schema("event")) +bucket = api.schema_model('Bucket', schema.get_json_schema("bucket")) +buckets_export = api.schema_model('Export', schema.get_json_schema("export")) # TODO: Construct all the models from JSONSchema? # A downside to contructing from JSONSchema: flask-restplus does not have marshalling support + info = api.model('Info', { 'hostname': fields.String(), 'version': fields.String(), 'testing': fields.Boolean(), }) -bucket = api.model('Bucket', { - 'id': fields.String(required=True, description='The buckets unique id'), - 'name': fields.String(required=False, description='The buckets readable and renameable name'), - 'type': fields.String(required=True, description='The buckets event type'), - 'client': fields.String(required=True, description='The name of the watcher client'), - 'hostname': fields.String(required=True, description='The hostname of the client that the bucket belongs to'), - 'created': fields.DateTime(required=True, description='The creation datetime of the bucket'), -}) - create_bucket = api.model('CreateBucket', { 'client': fields.String(required=True), 'type': fields.String(required=True), @@ -81,10 +78,11 @@ def format(self, value): }) query = api.model('Query', { - 'timeperiods': fields.List(fields.String, required=True, description='List of periods to query'), + 'timeperiods': fields.List(fields.String, required=True, description='List of periods to query'), 'query': fields.List(fields.String, required=True, description='String list of query statements'), }) + def copy_doc(api_method): """Decorator that copies another functions docstring to the decorated function. Used to copy the docstrings in ServerAPI over to the flask-restplus Resources. @@ -117,7 +115,7 @@ def get(self) -> Dict[str, Dict]: @api.route("/0/buckets/") class BucketResource(Resource): - @api.marshal_with(bucket) + @api.doc(model=bucket) @copy_doc(ServerAPI.get_bucket_metadata) def get(self, bucket_id): return app.api.get_bucket_metadata(bucket_id) @@ -254,21 +252,53 @@ def post(self): return {"type": type(qe).__name__, "message": str(qe)}, 400 -# EXPORTING +# EXPORT AND IMPORT + @api.route("/0/export") class ExportAllResource(Resource): + @api.doc(model=buckets_export) @copy_doc(ServerAPI.export_all) def get(self): - return app.api.export_all(), 200 + buckets_export = app.api.export_all() + payload = {"buckets": buckets_export} + response = make_response(json.dumps(payload)) + filename = "aw-buckets-export.json" + response.headers["Content-Disposition"] = "attachment; filename={}".format(filename) + return response # TODO: Perhaps we don't need this, could be done with a query argument to /0/export instead @api.route("/0/buckets//export") class BucketExportResource(Resource): + @api.doc(model=buckets_export) @copy_doc(ServerAPI.export_bucket) def get(self, bucket_id): - return app.api.export_bucket(bucket_id) + bucket_export = app.api.export_bucket(bucket_id) + payload = {"buckets": {bucket_export["id"]: bucket_export}} + response = make_response(json.dumps(payload)) + filename = "aw-bucket-export_{}.json".format(bucket_export["id"]) + response.headers["Content-Disposition"] = "attachment; filename={}".format(filename) + return response + + +@api.route("/0/import") +class ImportAllResource(Resource): + @api.expect(buckets_export) + @copy_doc(ServerAPI.import_all) + def post(self): + # If import comes from a form in th web-ui + if len(request.files) > 0: + # web-ui form only allows one file, but technically it's possible to + # upload multiple files at the same time + for filename, f in request.files.items(): + buckets = json.loads(f.stream.read())["buckets"] + app.api.import_all(buckets) + # Normal import from body + else: + buckets = request.get_json()["buckets"] + app.api.import_all(buckets) + return None, 200 # LOGGING