Commit 026817b0 authored by Alexandre Dulaunoy's avatar Alexandre Dulaunoy
Browse files

Ensure consistent JSON output with previous version of the API for datetime element.

The datetime bson output has a different format and use $date key when dealing
with datetime value. The output is just kept consistent with the previous version of the API.
parent 39a5fc7b
Loading
Loading
Loading
Loading
+13 −3
Original line number Diff line number Diff line
@@ -23,6 +23,8 @@ _runPath = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(_runPath, ".."))

from bson               import json_util
from bson.json_util import DEFAULT_JSON_OPTIONS
DEFAULT_JSON_OPTIONS.datetime_representation = 2
from dateutil.parser    import parse      as parse_datetime
from flask              import Flask, request, Response, render_template
from functools          import wraps
@@ -31,6 +33,7 @@ from redis import exceptions as redisExceptions
from tornado.httpserver import HTTPServer
from tornado.ioloop     import IOLoop
from tornado.wsgi       import WSGIContainer
import datetime

import lib.CVEs          as cves
import lib.DatabaseLayer as db
@@ -39,6 +42,12 @@ import lib.Toolkit as tk

from lib.Config import Configuration

def convertDatetime(dct=None):
  for key, val in dct.items():
    if isinstance(val, datetime.datetime):
      dct[key] = val.isoformat()
  return dct

class APIError(Exception):
  def __init__(self, message, status=500):
    self.message = message
@@ -74,6 +83,7 @@ class API():
  def addRoute(self, route):
    self.app.add_url_rule(route['r'], view_func=route['f'], methods=route['m'])


  #############
  # Decorator #
  #############
@@ -112,9 +122,9 @@ class API():
          if type(data) is not str:
            if type(data) is tuple:
              data = list(data)
              data[0] = json.dumps(data[0], indent=2, sort_keys=True, default=json_util.default)
              data[0] = json.dumps(convertDatetime(dct=data[0]), indent=4, sort_keys=True, default=json_util.default)
            else:
              data = (json.dumps(data, indent=2, sort_keys=True, default=json_util.default), 200)
              data = (json.dumps(convertDatetime(dct=data), indent=4, sort_keys=True, default=json_util.default), 200)
          return Response(data[0], mimetype=returnType), data[1]
      except:
        pass
@@ -272,7 +282,7 @@ class API():
    value=self.htmlDecode(value)
    regex = re.compile(re.escape(value), re.I)
    data = {'cves': db.via4Linked(key, regex)}
    cvssList=[float(x['cvss']) for x in data['cves'] if 'cvss' in x]
    cvssList=[float(x['cvss']) for x in data['cves'] if x.get('cvss')]
    if cvssList:
        data['stats']={'maxCVSS': max(cvssList), 'minCVSS': min(cvssList),'count':len(data['cves'])}
    else: