STRLCPY/tsunami-security-scanner

Add python plugin service Run rpc and tsunami plugin interface.

PiperOrigin-RevId: 464152054
Change-Id: I2e0c0c26fcef75beeb63efe83d85782bd1cd1ce6

John Y. Kim committed with Copybara-Service 2 years ago

2afba0ae

1 parent 816ec61f

Total 3 files Show one by one

■ ■ ■ ■ ■ ■

plugin_server/py/plugin_service.py

1	+	# Copyright 2022 Google LLC
2	+	#
3	+	# Licensed under the Apache License, Version 2.0 (the "License");
4	+	# you may not use this file except in compliance with the License.
5	+	# You may obtain a copy of the License at
6	+	#
7	+	# http://www.apache.org/licenses/LICENSE-2.0
8	+	#
9	+	# Unless required by applicable law or agreed to in writing, software
10	+	# distributed under the License is distributed on an "AS IS" BASIS,
11	+	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	+	# See the License for the specific language governing permissions and
13	+	# limitations under the License.
14	+	"""Python gRPC PluginService server adapter to provide communication with the Java client."""
15	+
16	+	from concurrent import futures
17	+	from typing import cast
18	+
19	+	from absl import logging
20	+	from tsunami.plugin_server.py import tsunami_plugin
21	+	from tsunami.proto import detection_pb2
22	+	from tsunami.proto import plugin_representation_pb2
23	+	from tsunami.proto import plugin_service_pb2
24	+	from tsunami.proto import plugin_service_pb2_grpc
25	+
26	+	RunResponse = plugin_service_pb2.RunResponse
27	+	_PluginType = plugin_representation_pb2.PluginInfo.PluginType
28	+
29	+	_DETECTION_TIMEOUT = 60
30	+
31	+
32	+	class PluginServiceServicer(plugin_service_pb2_grpc.PluginServiceServicer):
33	+	"""PluginService server implementation for communication with the Java client.
34	+
35	+	This class executes requests called by the Java client. All request types are
36	+	given by the plugin_service proto definition.
37	+
38	+	"""
39	+
40	+	def __init__(self, py_plugins: list[tsunami_plugin.TsunamiPlugin],
41	+	max_workers: int):
42	+	self.py_plugins = py_plugins
43	+	self.max_workers = max_workers
44	+
45	+	def Run(
46	+	self, request: plugin_service_pb2.RunRequest,
47	+	servicer_context: plugin_service_pb2_grpc.PluginServiceServicer
48	+	) -> RunResponse:
49	+	logging.info('Received Run request = %s', request)
50	+	report_list = detection_pb2.DetectionReportList()
51	+
52	+	detection_futures = []
53	+
54	+	with futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
55	+	for matched_plugin in request.plugins:
56	+	plugin_def = matched_plugin.plugin
57	+	for plugin in self.py_plugins:
58	+	if plugin.GetPluginDefinition() == plugin_def:
59	+	logging.info('Running python plugin %s.', type(plugin).__name__)
60	+
61	+	if plugin_def.info.type is _PluginType.VULN_DETECTION:
62	+	plugin = cast(tsunami_plugin.VulnDetector, plugin)
63	+	detection_futures.append(
64	+	executor.submit(plugin.Detect, request.target,
65	+	matched_plugin.services))
66	+
67	+	for detection in detection_futures:
68	+	report_list.detection_reports.extend(
69	+	detection.result(timeout=_DETECTION_TIMEOUT).detection_reports)
70	+
71	+	response = RunResponse()
72	+	response.reports.CopyFrom(report_list)
73	+	return response
74	+

■ ■ ■ ■ ■ ■

plugin_server/py/plugin_service_test.py

1	+	# Copyright 2022 Google LLC
2	+	#
3	+	# Licensed under the Apache License, Version 2.0 (the "License");
4	+	# you may not use this file except in compliance with the License.
5	+	# You may obtain a copy of the License at
6	+	#
7	+	# http://www.apache.org/licenses/LICENSE-2.0
8	+	#
9	+	# Unless required by applicable law or agreed to in writing, software
10	+	# distributed under the License is distributed on an "AS IS" BASIS,
11	+	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	+	# See the License for the specific language governing permissions and
13	+	# limitations under the License.
14	+	"""Tests for plugin_service."""
15	+
16	+	import time
17	+
18	+	import grpc_testing
19	+	import ipaddr
20	+
21	+	from google.protobuf import timestamp_pb2
22	+	from testing.pybase import googletest
23	+	from tsunami.plugin_server.py import plugin_service
24	+	from tsunami.plugin_server.py import tsunami_plugin
25	+	from tsunami.proto import detection_pb2
26	+	from tsunami.proto import network_pb2
27	+	from tsunami.proto import network_service_pb2
28	+	from tsunami.proto import plugin_representation_pb2
29	+	from tsunami.proto import plugin_service_pb2
30	+	from tsunami.proto import reconnaissance_pb2
31	+	from tsunami.proto import vulnerability_pb2
32	+
33	+	_NetworkEndpoint = network_pb2.NetworkEndpoint
34	+	_NetworkService = network_service_pb2.NetworkService
35	+	_PluginInfo = plugin_representation_pb2.PluginInfo
36	+	_TargetInfo = reconnaissance_pb2.TargetInfo
37	+	_AddressFamily = network_pb2.AddressFamily
38	+	_ServiceDescriptor = plugin_service_pb2.DESCRIPTOR.services_by_name[
39	+	'PluginService']
40	+	_RunMethod = _ServiceDescriptor.methods_by_name['Run']
41	+	MAX_WORKERS = 1
42	+
43	+
44	+	class PluginServiceTest(googletest.TestCase):
45	+
46	+	@classmethod
47	+	def setUpClass(cls):
48	+	super().setUpClass()
49	+	cls.test_plugin = FakeVulnDetector()
50	+	cls._time = grpc_testing.strict_fake_time(time.time())
51	+	cls._server = grpc_testing.server_from_dictionary(
52	+	{
53	+	_ServiceDescriptor:
54	+	plugin_service.PluginServiceServicer(
55	+	py_plugins=[cls.test_plugin], max_workers=MAX_WORKERS),
56	+	}, cls._time)
57	+
58	+	cls._channel = grpc_testing.channel(
59	+	plugin_service_pb2.DESCRIPTOR.services_by_name.values(), cls._time)
60	+
61	+	@classmethod
62	+	def tearDownClass(cls):
63	+	cls._channel.close()
64	+	super().tearDownClass()
65	+
66	+	def test_run_plugins_registered_returns_valid_response(self):
67	+	plugin_to_test = FakeVulnDetector()
68	+	endpoint = _build_network_endpoint('1.1.1.1', 80)
69	+	service = _NetworkService(
70	+	network_endpoint=endpoint,
71	+	transport_protocol=network_pb2.TCP,
72	+	service_name='http')
73	+	target = _TargetInfo(network_endpoints=[endpoint])
74	+	services = [service]
75	+	request = plugin_service_pb2.RunRequest(
76	+	target=target,
77	+	plugins=[
78	+	plugin_service_pb2.MatchedPlugin(
79	+	services=services, plugin=plugin_to_test.GetPluginDefinition())
80	+	])
81	+
82	+	rpc = self._server.invoke_unary_unary(_RunMethod, (), request, None)
83	+	response, _, _, _ = rpc.termination()
84	+
85	+	self.assertLen(response.reports.detection_reports, 1)
86	+	self.assertEqual(
87	+	plugin_to_test._BuildFakeDetectionReport(
88	+	target=target, network_service=services[0]),
89	+	response.reports.detection_reports[0])
90	+
91	+	def test_run_no_plugins_registered_returns_empty_response(self):
92	+	endpoint = _build_network_endpoint('1.1.1.1', 80)
93	+	target = _TargetInfo(network_endpoints=[endpoint])
94	+	request = plugin_service_pb2.RunRequest(target=target, plugins=[])
95	+
96	+	rpc = self._server.invoke_unary_unary(_RunMethod, (), request, None)
97	+	response, _, _, _ = rpc.termination()
98	+
99	+	self.assertEmpty(response.reports.detection_reports)
100	+
101	+
102	+	def _build_network_endpoint(ip: str, port: int) -> _NetworkEndpoint:
103	+	return _NetworkEndpoint(
104	+	type=_NetworkEndpoint.IP,
105	+	ip_address=network_pb2.IpAddress(address_family=_get_address_family(ip)),
106	+	port=network_pb2.Port(port_number=port))
107	+
108	+
109	+	def _get_address_family(ip: str) -> _AddressFamily:
110	+	inet_addr = ipaddr.IPAddress(ip)
111	+	if inet_addr.version == 4:
112	+	return _AddressFamily.IPV4
113	+	elif inet_addr.version == 6:
114	+	return _AddressFamily.IPV6
115	+	else:
116	+	raise ValueError('Unknown IP address family for IP \'%s\'' % ip)
117	+
118	+
119	+	class FakeVulnDetector(tsunami_plugin.VulnDetector):
120	+	"""Fake Vulnerability detector class for testing only."""
121	+
122	+	def GetPluginDefinition(self):
123	+	return tsunami_plugin.PluginDefinition(
124	+	info=_PluginInfo(
125	+	type=_PluginInfo.VULN_DETECTION,
126	+	name='fake',
127	+	version='v0.1',
128	+	description='fake description',
129	+	author='fake author'),
130	+	target_service_name=plugin_representation_pb2.TargetServiceName(
131	+	value=['fake service']),
132	+	target_software=plugin_representation_pb2.TargetSoftware(
133	+	name='fake software'),
134	+	for_web_service=False)
135	+
136	+	def Detect(self, target, matched_services):
137	+	return detection_pb2.DetectionReportList(detection_reports=[
138	+	self._BuildFakeDetectionReport(target, matched_services[0])
139	+	])
140	+
141	+	def _BuildFakeDetectionReport(self, target, network_service):
142	+	return detection_pb2.DetectionReport(
143	+	target_info=target,
144	+	network_service=network_service,
145	+	detection_timestamp=timestamp_pb2.Timestamp(nanos=1234567890),
146	+	detection_status=detection_pb2.VULNERABILITY_VERIFIED,
147	+	vulnerability=vulnerability_pb2.Vulnerability(
148	+	main_id=vulnerability_pb2.VulnerabilityId(
149	+	publisher='GOOGLE', value='FakeVuln1'),
150	+	severity=vulnerability_pb2.CRITICAL,
151	+	title='FakeTitle1',
152	+	description='FakeDescription1'))
153	+
154	+
155	+	# TODO(b/239628051): Add a failed VulnDetector class to test failed cases.
156	+
157	+	if __name__ == '__main__':
158	+	googletest.main()
159	+

■ ■ ■ ■ ■ ■

plugin_server/py/tsunami_plugin.py

1	+	# Copyright 2022 Google LLC
2	+	#
3	+	# Licensed under the Apache License, Version 2.0 (the "License");
4	+	# you may not use this file except in compliance with the License.
5	+	# You may obtain a copy of the License at
6	+	#
7	+	# http://www.apache.org/licenses/LICENSE-2.0
8	+	#
9	+	# Unless required by applicable law or agreed to in writing, software
10	+	# distributed under the License is distributed on an "AS IS" BASIS,
11	+	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	+	# See the License for the specific language governing permissions and
13	+	# limitations under the License.
14	+	"""Interface that all Python TsunamiPlugins will need to implement to run detection."""
15	+	import abc
16	+
17	+	from tsunami.proto import detection_pb2
18	+	from tsunami.proto import network_service_pb2
19	+	from tsunami.proto import plugin_representation_pb2
20	+	from tsunami.proto import reconnaissance_pb2
21	+
22	+	_TargetInfo = reconnaissance_pb2.TargetInfo
23	+	_NetworkService = network_service_pb2.NetworkService
24	+	_DetectionReportList = detection_pb2.DetectionReportList
25	+	PluginDefinition = plugin_representation_pb2.PluginDefinition
26	+
27	+
28	+	class TsunamiPlugin(metaclass=abc.ABCMeta):
29	+
30	+	@abc.abstractmethod
31	+	def GetPluginDefinition(self) -> PluginDefinition:
32	+	pass
33	+
34	+	@classmethod
35	+	def __subclasshook__(cls, subclass: abc.ABCMeta) -> bool:
36	+	return (hasattr(subclass, 'GetPluginDefinition') and
37	+	callable(subclass.GetPluginDefinition))
38	+
39	+
40	+	class VulnDetector(TsunamiPlugin):
41	+	"""A TsunamiPlugin that detects potential vulnerabilities on the target.
42	+
43	+	Usually a vulnerability detector takes the information about an exposed
44	+	network service, detects whether the service is vulnerable to a
45	+	specific vulnerability, and reports the detection results.
46	+	"""
47	+
48	+	@abc.abstractmethod
49	+	def Detect(self, target: _TargetInfo,
50	+	matched_services: list[_NetworkService]) -> _DetectionReportList:
51	+	"""Run detection logic for the target.
52	+
53	+	Args:
54	+	target: Information abouut the scanning target itself.
55	+	matched_services: A list of network services whose vulnerabilities could
56	+	be detected by this plugin.
57	+
58	+	Returns:
59	+	A DetectionReportList for all the vulnerabilities of the scanning target.
60	+	"""
61	+