v0.6.1 - code cleanup, better session handling, tests added

README.md

@@ -112,7 +112,7 @@ CWE (Common Weakness Enumeration), ... as well as additional information (RedHat
 | cveid               | CVE number            | `CVE-2014-0160`          |
 
 ```python
->>> cve.cve('CVE-2014-0160')
+>>> cve.id('CVE-2014-0160')
 ```
 
 <br/>
@@ -189,15 +189,17 @@ Outputs the last `n` amount of vulnerabilities. If the limit is not specified, t
 
 ##### Description
 
+**DISABLED ON cve.circl.lu**
+
 Returns all CVEs that are linked by a given key/value pair.
 
 | Argument            | Description                         | Example                     |
 | :-------------------| :---------------------------------- | :-------------------------- |
-| key                 | The key to link CVEs on             | `msbulletin.bulletin_id`    |
+| key                 | The key to link CVEs on             | `refmap.ms`                 |
 | value               | The value for the given key         | `MS16-098`                  |
 
 ```python
->>> cve.link('msbulletin.bulletin_id/MS16-098')
+>>> cve.link('refmap.ms/MS16-098')
 ```
 
 <br/>

ares/__init__.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 __title__ = 'ares'
-__version__ = '0.6.0'
+__version__ = '0.6.1'
 __author__ = 'Martin Simon <me@martinsimon.me>'
 __repo__ = 'https://github.com/barnumbirr/ares'
 __license__ = 'Apache v2.0 License'

ares/core.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-import json
 import requests
 
 class CVESearch(object):
@@ -13,74 +12,58 @@ class CVESearch(object):
 	def __init__(self, base_url = __DEFAULT_BASE_URL, request_timeout = __DEFAULT_TIMEOUT):
 		self.base_url = base_url
 		self.request_timeout = request_timeout
+		self.session = self._create_session()
 
-	@property
-	def session(self):
-		if not self._session:
-			self._session = requests.Session()
-			self._session.headers.update({'Content-Type': 'application/json'})
-			self._session.headers.update({'User-agent': 'ares - python wrapper \
-		around cve.circl.lu (github.com/barnumbirr/ares)'})
-		return self._session
+	@staticmethod
+	def _create_session():
+		session = requests.Session()
+		user_agent = 'ares - python wrapper around cve.circl.lu (github.com/barnumbirr/ares)'
+		session.headers.update({'Content-Type': 'application/json'})
+		session.headers.update({'User-agent': user_agent})
+		return session
 
 	def __request(self, endpoint, query):
 		# There is probably a more elegant way to do this ¯\_(ツ)_/¯
 		if query:
-			response_object = self.session.get(requests.compat.urljoin(self.base_url, endpoint + query),
-		                                   timeout = self.request_timeout)
+			response = self.session.get(requests.compat.urljoin(self.base_url, endpoint + query),
+		                                timeout = self.request_timeout)
 		else:
-			response_object = self.session.get(requests.compat.urljoin(self.base_url, endpoint),
-		                                   timeout = self.request_timeout)
+			response = self.session.get(requests.compat.urljoin(self.base_url, endpoint),
+		                                timeout = self.request_timeout)
 
-		try:
-			response = json.loads(response_object.text)
-		except Exception as e:
-			return e
-
-		return response
+		response.raise_for_status()
+		return response.json()
 
 	def browse(self, param=None):
-		response = self.__request('browse/', query=param)
-		return response
+		return self.__request('browse/', query=param)
 
 	def capec(self, param):
-		response = self.__request('capec/', query=param)
-		return response
+		return self.__request('capec/', query=param)
 
 	# def cpe22(self, param):
-	# 	response = self.__request('cpe2.2/', query=param)
-	# 	return response
-
+	# 	return self.__request('cpe2.2/', query=param)
 
 	# def cpe23(self, param):
-	# 	response = self.__request('cpe2.3/', query=param)
-	# 	return response
-
-	def cve(self, param):
-		response = self.__request('cve/', query=param)
-		return response
+	# 	return self.__request('cpe2.3/', query=param)
 
 	# def cvefor(self, param):
-	# 	response = self.__request('cvefor/', query=param)
-	# 	return response
+	# 	return self.__request('cvefor/', query=param)
 
 	def cwe(self):
 		""" Outputs a list of all CWEs (Common Weakness Enumeration). """
-		response = self.__request('cwe', query=None)
-		return response
+		return self.__request('cwe', query=None)
 
 	def dbinfo(self):
-		response = self.__request('dbInfo', query=None)
-		return response
+		return self.__request('dbInfo', query=None)
+
+	def id(self, param):
+		return self.__request('cve/', query=param)
 
-	def last(self, param):
-		response = self.__request('last/', query=param)
-		return response
+	def last(self, param=None):
+		return self.__request('last/', query=param)
 
-	def link(self, param):
-		response = self.__request('link/', query=param)
-		return response
+	# def link(self, param):
+	# 	return self.__request('link/', query=param)
 
 	# def search(self, param):
-	# 	response = self.__request('search/', query=param)
-	# 	return response
+	# 	return self.__request('search/', query=param)

tests/__init__.py

@@ -0,0 +1,2 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-

tests/test_ares.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import unittest
+try:
+    from .core import CVESearch
+except ModuleNotFoundError:
+    from ares import CVESearch
+
+class TestCVEAPI(unittest.TestCase):
+
+    def setUp(self):
+        self.cve = CVESearch()
+
+    def tearDown(self):
+        self.cve.session.close()
+
+    def test_init(self):
+        self.assertTrue(isinstance(self.cve, CVESearch))
+
+    def test_session_headers(self):
+        user_agent = 'ares - python wrapper around cve.circl.lu (github.com/barnumbirr/ares)'
+        self.assertEqual(self.cve.session.headers["Content-Type"], "application/json")
+        self.assertEqual(self.cve.session.headers["User-agent"], user_agent)
+
+    @unittest.skip("Test too aggressive for provider.")
+    def test_empty_browse(self):
+        response = self.cve.browse()
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+        self.assertIsNone(response["product"])
+        self.assertIsInstance(response["vendor"], list)
+        self.assertTrue(len(response["vendor"]) > 1000)
+
+    def test_browse(self):
+        response = self.cve.browse(param="python-requests")
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+        self.assertEqual(response["vendor"], "python-requests")
+
+    def test_capec(self):
+        response = self.cve.capec(param="13")
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+        self.assertEqual(response["name"], "Subverting Environment Variable Values")
+
+    @unittest.skip("Endpoint disabled on cve.circl.lu")
+    def test_cpe22(self):
+        response = self.cve.cpe22('cpe:2.3:o:microsoft:windows_vista:6.0:sp1:-:-:home_premium:-:-:x64:-')
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, str)
+        self.assertEqual(response, "cpe:/o:microsoft:windows_vista:6.0:sp1:~~home_premium~~x64~")
+
+    @unittest.skip("Endpoint disabled on cve.circl.lu")
+    def test_cpe23(self):
+        response = self.cve.cpe23('cpe:/o:microsoft:windows_vista:6.0:sp1:~-~home_premium~-~x64~-')
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, str)
+        self.assertEqual(response, "cpe:2.3:o:microsoft:windows_vista:6.0:sp1:-:-:home_premium:-:-:x64")
+
+    @unittest.skip("Endpoint disabled on cve.circl.lu")
+    def test_cvefor(self):
+        response = self.cve.cvefor('cpe:/o:microsoft:windows_vista:6.0:sp1:~-~home_premium~-~x64~-')
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+        self.assertEqual(response["id"], "CVE-2005-0100")
+
+    @unittest.skip("Test too aggressive for provider.")
+    def test_cwe(self):
+        response = self.cve.cwe()
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+
+    def test_db_info(self):
+        response = self.cve.dbinfo()
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+
+    def test_id(self):
+        response = self.cve.id(param="CVE-2015-2296")
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+        self.assertEqual(response["Published"], "2015-03-18T16:59:00")
+
+    def test_bad_id(self):
+        response = self.cve.id(param="CVE-not-real")
+        self.assertIsNone(response)
+
+    def test_last(self):
+        response = self.cve.last()
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, list)
+        self.assertEqual(len(response), 30)
+
+    @unittest.skip("Endpoint disabled on cve.circl.lu")
+    def test_link(self):
+        response = self.cve.link(param="refmap.ms/CVE-2016-3309")
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+        self.assertEqual(response["cves"]["cwe"], "CWE-264")
+
+    @unittest.skip("Endpoint disabled on cve.circl.lu")
+    def test_search_vendor(self):
+        response = self.cve.search(param="python-requests")
+        self.assertIsNotNone(response)
+        self.assertIsInstance(response, dict)
+        self.assertIsInstance(response["data"], list)
+
+if __name__ == "__main__":
+    unittest.main()