diff --git a/README.md b/README.md index 1128031c6..d78ad5c84 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ # VMware Carbon Black Cloud Python SDK -**Latest Version:** 1.3.1 +**Latest Version:** 1.3.2
-**Release Date:** 15 June 2021 +**Release Date:** 10 August 2021 [![Coverage Status](https://coveralls.io/repos/github/carbonblack/carbon-black-cloud-sdk-python/badge.svg?t=Id6Baf)](https://coveralls.io/github/carbonblack/carbon-black-cloud-sdk-python) [![Codeship Status for carbonblack/carbon-black-cloud-sdk-python](https://app.codeship.com/projects/9e55a370-a772-0138-aae4-129773225755/status?branch=develop)](https://app.codeship.com/projects/402767) diff --git a/VERSION b/VERSION index 3a3cd8cc8..1892b9267 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.3.1 +1.3.2 diff --git a/docs/changelog.rst b/docs/changelog.rst index 270ddfd2b..f2324f1cd 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -1,5 +1,23 @@ Changelog ================================ +CBC SDK 1.3.2 - Released August 10, 2021 +-------------------------------- + +New Features: + +* Added asynchronous query options to Live Response APIs. +* Added functionality for Watchlists, Reports, and Feeds to simplify developer interaction. + +Updates: + +* Added documentation on the mapping between permissions and Live Response commands. + +Bug Fixes: + +* Fixed an error using the STIX/TAXII example with Cabby. +* Fixed a potential infinite loop in getting detailed search results for enriched events and processes. +* Comparison now case-insensitive on UBS download. + CBC SDK 1.3.1 - Released June 15, 2021 -------------------------------- diff --git a/docs/conf.py b/docs/conf.py index 6adbdcc56..d059e0665 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -23,7 +23,7 @@ author = 'Developer Relations' # The full version, including alpha/beta/rc tags -release = '1.3.1' +release = '1.3.2' # -- General configuration --------------------------------------------------- diff --git a/docs/guides-and-resources.rst b/docs/guides-and-resources.rst index a45568573..38297593e 100755 --- a/docs/guides-and-resources.rst +++ b/docs/guides-and-resources.rst @@ -20,6 +20,7 @@ Guides * :doc:`live-response` - Live Response allows security operators to collect information and take action on remote endpoints in real time. * :doc:`unified-binary-store` - The unified binary store (UBS) is responsible for storing all binaries and corresponding metadata for those binaries. * :doc:`users-grants` - Work with users and access grants. +* :doc:`watchlists-feeds-reports` - Work with Enterprise EDR watchlists, feeds, reports, and Indicators of Compromise (IOCs). Examples -------- diff --git a/docs/live-response.rst b/docs/live-response.rst index bac49a5d4..fc13fdf4e 100755 --- a/docs/live-response.rst +++ b/docs/live-response.rst @@ -8,6 +8,77 @@ You can use Live Response with the Carbon Black Cloud Python SDK to: * Dump contents of physical memory * Execute, terminate and list processes +Before any commands are sent to the live response session, the proper permissions need to be configured for the Custom Key that is used. +The below table explains what permissions are needed for each of the SDK commands. + ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| Command | Required Permissions | Explanation | ++===================================================+========================================================+======================================================+ +| | Create LR session for device | **CREATE**, **READ** org.liveresponse.session | CREATE is needed to start the LR session and | +| | device.lr_session() | | READ is needed to check the status of the command | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Close session | **READ**, **DELETE** org.liveresponse.session | DELETE is needed to terminate the LR session and | +| | lr_session.close() | | READ is needed to check the status of the command | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Get Raw File | **READ** org.liveresponse.file | | +| | lr_session.get_raw_file(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Get File | **READ** org.liveresponse.file | | +| | lr_session.get_file(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Upload File | **CREATE**, **READ** org.liveresponse.file | CREATE is needed to upload the file and READ is | +| | lr_session.put_file(...) | | needed to check the status of the command | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Delete file | **READ**, **DELETE** org.liveresponse.file | DELETE is needed to delete the file and READ is | +| | lr_session.delete_file(...) | | needed to check the status of the command | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | List Directory | **READ** org.liveresponse.file | | +| | lr_session.list_directory(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Create Directory | **CREATE**, **READ** org.liveresponse.file | CREATE is needed to create the directory and | +| | lr_session.create_directory(...) | | READ is needed to check the status of the command | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Walk Directory | **READ** org.liveresponse.file | | +| | lr_session.walk(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Kill Process | **READ**, **DELETE** org.liveresponse.process | DELETE is needed to kill the process and READ is | +| | lr_session.kill_process(...) | | needed to check the status of the command | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Create Process | | **EXECUTE** org.liveresponse.process | If wait_for_completion = False, wait_for_output = | +| | lr_session.create_process(...) | | OR | False only EXECUTE is needed. | +| | | **EXECUTE** org.liveresponse.process | Otherwise also file permissions are needed. | +| | | **READ**, **DELETE** org.liveresponse.file | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | List Processes | **READ** org.liveresponse.process | | +| | lr_session.list_processes(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | List Registry Keys and Values | **READ** org.liveresponse.registry | | +| | lr_session.list_registry_keys_and_values(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | List Registry Values | **READ** org.liveresponse.registry | | +| | lr_session.list_registry_values(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Get Registry Value | **READ** org.liveresponse.registry | | +| | lr_session.get_registry_value(...) | | | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Set Registry | **READ**, **UPDATE** org.liveresponse.registry | UPDATE is needed to set/create the value for the | +| | lr_session.set_registry_value(...) | | registry and READ to check the status of the command | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Create Registry Key | **CREATE**, **READ** org.liveresponse.registry | CREATE is needed to create the key and READ to | +| | lr_session.create_registry_key(...) | | check the status of the command. | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Delete Registry Key | **READ**, **DELETE** org.liveresponse.registry | DELETE is needed to delete the key and READ to | +| | lr_session.delete_registry_key(...) | | check the status of the command. | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Delete Registry Value | **READ**, **DELETE** org.liveresponse.registry | DELETE is needed to delete the value and READ to | +| | lr_session.delete_registry_value(...) | | check the status of the command. | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ +| | Memdump | **READ** org.liveresponse.memdump | The command to dump the memory includes three | +| | lr_session.memdump(...) | **READ**, **DELETE** org.liveresponse.file | commands - dumping the memory in a file on the | +| | | remote machine, downloading the file on the local | +| | | machine and deleting the file. | ++---------------------------------------------------+--------------------------------------------------------+------------------------------------------------------+ + To send commands to an endpoint, first establish a "session" with a device. Establish A Session With A Device diff --git a/docs/unified-binary-store.rst b/docs/unified-binary-store.rst index 27b61f9d1..2da885bb4 100755 --- a/docs/unified-binary-store.rst +++ b/docs/unified-binary-store.rst @@ -10,7 +10,7 @@ Get Download URL :: >>> from cbc_sdk import CBCloudAPI - >>> api = CBCloudAPI(profile='sample') + >>> cb = CBCloudAPI(profile='sample') >>> from cbc_sdk.enterprise_edr.ubs import Binary >>> sha256_hash = '8005557c1614c1e2c89f7db3702199de2b1e4605718fa32ff6ffdb2b41ed3759' >>> binary = Binary(cb, sha256_hash) @@ -29,7 +29,7 @@ We could set expiration period for the download link (in seconds). :: >>> from cbc_sdk import CBCloudAPI - >>> api = CBCloudAPI(profile='sample') + >>> cb = CBCloudAPI(profile='sample') >>> from cbc_sdk.enterprise_edr.ubs import Binary >>> sha256_hash = '8005557c1614c1e2c89f7db3702199de2b1e4605718fa32ff6ffdb2b41ed3759' >>> binary = Binary(cb, sha256_hash) @@ -48,7 +48,7 @@ Currently querying binaries is not possible, but we could use the following synt :: >>> from cbc_sdk import CBCloudAPI - >>> api = CBCloudAPI(profile='sample') + >>> cb = CBCloudAPI(profile='sample') >>> from cbc_sdk.enterprise_edr.ubs import Binary >>> sha256_hash = '8005557c1614c1e2c89f7db3702199de2b1e4605718fa32ff6ffdb2b41ed3759' >>> binary = cb.select(Binary, sha256_hash) @@ -56,8 +56,7 @@ Currently querying binaries is not possible, but we could use the following synt ... https://cdc-file-storage-staging-us-east-1.s3.amazonaws.com/80/05/55/7c/16/14/c1/ -*Note: If we try to use* :code:`binary = cb.select(Binary)` *, it will fail with exception that the model is non queryable model.* +*Note: If we try to use* :code:`binary = cb.select(Binary)` *, it will fail with exception that the model is non queryable model.* Find the full documentation at `Unified Binary Store `_. - diff --git a/docs/watchlists-feeds-reports.rst b/docs/watchlists-feeds-reports.rst new file mode 100644 index 000000000..4e9eac2bf --- /dev/null +++ b/docs/watchlists-feeds-reports.rst @@ -0,0 +1,269 @@ +Watchlists, Feeds, Reports, and IOCs +==================================== +Watchlists are a powerful feature of Carbon Black Cloud Enterprise EDR. They allow an organization to set-and-forget +searches on their endpoints' incoming events data, providing the administrator the opportunity to sift through high +volumes of activity and focus attention on those that matter. + +**Note:** Use of these APIs requires that the organization be enabled for Enterprise EDR. Verify this by logging into +the Carbon Black Cloud Console, opening the menu in the upper right corner, and checking for an ``ENABLED`` flag +against the "Enterprise EDR" entry. + +All examples here assume that a Carbon Black Cloud SDK connection has been set up, such as with the following code: + +:: + + >>> from cbc_sdk import CBCloudAPI + >>> api = CBCloudAPI(profile='sample') + +Setting up a connection is documented here: :doc:`getting-started` + +About the Objects +----------------- +An *indicator of compromise* (IOC) is a query, list of strings, or list of regular expressions which constitutes +actionable threat intelligence that the Carbon Black Cloud is set up to watch for. Any activity that matches one of +these may indicate a compromise of an endpoint. + +A *report* groups one or more IOCs together, which may reflect a number of possible conditions to look for, or a number +of conditions related to a particular target program or type of malware. Reports can be used to organize IOCs. + +A *watchlist* contains reports (either directly or through a feed) that the Carbon Black Cloud is matching against +events coming from the endpoints. A positive match will trigger a "hit," which may be logged or result in an alert. + +A *feed* contains reports which have been gathered by a single source. They resemble "potential watchlists." +A watchlist may be easily subscribed to a feed, so that any reports in the feed act as if they were in the watchlist +itself, triggering logs or alerts as appropriate. + +Setting Up a Basic Custom Watchlist +----------------------------------- +Creating a custom watchlist that can watch incoming events and/or generate alerts requires three steps: + +1. Create a report including one or more Indicators of Compromise (IOCs). +2. Add that report to a watchlist. +3. Enable alerting on the watchlist. + +Creating a Report ++++++++++++++++++ +In this example, a report is created, adding one or more IOCs to it: + +:: + + >>> from cbc_sdk.enterprise_edr import Report, IOC_V2 + >>> builder = Report.create(api, "Unsigned Browsers", "Unsigned processes impersonating browsers", 5) + >>> builder.add_tag("compliance").add_tag("unsigned_browsers") + >>> builder.add_ioc(IOC_V2.create_query(api, "unsigned-chrome", + ... "process_name:chrome.exe NOT process_publisher_state:FILE_SIGNATURE_STATE_SIGNED")) + >>> report = builder.build() + >>> report.save_watchlist() + +Reports should always be given a ``title`` that's sufficiently unique within your organization, so as to minimize +the chances of confusing two or more Reports with each other. Carbon Black Cloud will generate unique ``id`` values +for each report, but does not enforce any uniqueness constraint on the ``title`` of reports. + +Alternatively, you can update an existing report, adding more IOCs and/or replacing existing ones. To find an existing +report associated with a watchlist, you must look in the watchlist's ``reports`` collection: + +:: + + >>> from cbc_sdk.enterprise_edr import Watchlist, Report, IOC_V2 + >>> watchlist = api.select('Watchlist', 'R4cMgFIhRaakgk749MRr6Q') + >>> report_list = [report for report in watchlist.reports where report.id == '47474d40-1f94-4995-b6d9-1d1eea3528b3'] + >>> report = report_list[0] + >>> report.append_iocs([IOC_V2.create_query(api, 'evil-connect', 'netconn_ipv4:10.8.16.4')]) + >>> report.save_watchlist() + +Adding the Report to a Watchlist +++++++++++++++++++++++++++++++++ +Now, add the new Report to a new Watchlist: + +:: + + >>> from cbc_sdk.enterprise_edr import Watchlist + >>> builder = Watchlist.create(api, "Suspicious Applications") + >>> builder.set_description("Any signs of suspicious applications running on endpoints").add_reports([report]) + >>> watchlist = builder.build() + >>> watchlist.save() + +If you already have an existing Watchlist you wish to enhance, you can add Reports to the existing Watchlist: + +:: + + >>> # "report" contains the Report that was created in the previous example + >>> from cbc_sdk.enterprise_edr import Watchlist + >>> watchlist = api.select('Watchlist', 'R4cMgFIhRaakgk749MRr6Q') + >>> watchlist.add_reports([report]) + +Enabling Alerting on a Watchlist +++++++++++++++++++++++++++++++++ +When either the ``alerts_enabled`` or ``tags_enabled`` attributes of a watchlist are ``True``, that Watchlist will +create data you can act on - either alerts or hits, respectively; if both are ``False``, the Watchlist is effectively +disabled. + +Once you have the Watchlist configured with the IOCs that are generating the kinds of hits (results) you are after, +you can enable Alerting for the Watchlist, which will allow matches against the reports in the watchlist to generate +alerts. If a watchlist identifies suspicious behavior and known threats in your environment, you will want to enable +alerts to advise you of situations where you may need to take action or modify policies. + +:: + + >>> watchlist.enable_alerts() + +A Closer Look at IOCs +--------------------- +In this document, only the "v2" IOCs are covered; the "v1" IOCs are only provided for backwards compatibility +reasons. They are officially deprecated, and are converted, internally, to this type. + +IOCs can be classified into two general types, depending on their ``match_type`` value: + +*Query IOCs* are those with a ``match_type`` of ``query``; their ``values_list`` contains a single string that +specifies a query compatible with process searches. For example, the following IOC looks for the process ``git.exe`` +that does *not* connect to one of a specified list of IP addresses: + +:: + + { + "id": "example_1", + "match_type": "query", + "values": ["process_name:git.exe NOT (netconn_ipv4:35.158.151.206 OR netconn_ipv4:1.1.244.78 + OR netconn_ipv4:80.18.61.229 OR netconn_ipv4:80.18.61.228)"] + } + +Query IOCs must always use field-prefixed queries (key-value pairs); they do not support just searching for a value +without a field specified. Values in query clauses that do not specify fields will be ignored. + +:Wrong: ``process_name:chrome.exe AND 192.168.1.1`` +:Right: ``process_name:chrome.exe AND netconn_ipv4:192.168.1.1`` + +Query IOCs may search on CIDR address ranges, e.g.: ``netconn_ipv4:192.168.0.0/16``. + +Query IOCs are searched every 5 minutes by the Carbon Black Cloud, and are tested against a rolling window of the +last hour's worth of data for the organization. (They will *not* generate hits or alerts for process attributes that +were reported more than an hour in the past.) They may employ any searchable field as documented +`here `_, +and may employ complex query logic. + +*Ingress IOCs* are those with a ``match_type`` of ``equality`` or ``regex``; they use the ``field`` element to specify +the name of a field to examine the value of, and the ``values_list`` element to specify a list of values to match +against (in the case of ``match_type`` being ``equality``) or regular expressions to match against (in the case of +``match_type`` being ``regex``). For example, this IOC will match any process that initiates a connection to one of +two listed IP addresses: + +:: + + { + "id": "example_2", + "match_type": "equality", + "field": "netconn_ipv4", + "values": ["8.8.8.8", "1.160.120.15"] + } + +This IOC will match any process running with an executable name beginning with "quake": + +:: + + { + "id": "example_3", + "match_type": "regex", + "field": "process_name", + "values": ["quake.*\\.exe"] + } + +(Note the use of the backslash to escape the '.' that separates the file extension from the name. It must be doubled +to escape it in Python itself.) + +Ingress IOCs are searched as soon as the data is received from any endpoint, and may use any process field +(as documented +`here `_; +the fields that may be used in this context are tagged with ``PROCESS``) +in their ``field`` element, whether searchable or not. For the searches they are capable of, they are more efficient +than query IOCs, and also easier to add additional search target values to. They can, however, only search on a single +field at a time. + +**Note:** Ingress IOCs cannot be edited in the Carbon Black Cloud console UI at this time, due to a UI limitation +on editing two properties of an IOC at the same time. + +You *can* include more than one entry (query or match element) in an individual IOC, but in order to ignore or disable +one of those entries, you would either have to edit the IOC or disable it entirely (thus disabling *all* entries in +that IOC). It is recommended to use only one entry per IOC, for ease of management, unless you have already vetted the +entries and don't expect to have to disable them individually. + +Both IOCs and reports may include a ``link`` property, which is used by the Carbon Black Cloud console UI as a hint +to indicate that this IOC or report is being managed outside of the console. If this property is not ``None``, +the console UI will disable the ability to edit the IOC or report, but they can still be edited via the API. + +Tips for Using IOCs ++++++++++++++++++++ +* You can safely ignore certain fields in an IOC. For example, fields like ``alert_id`` and ``process_guid`` will + always uniquely identify just a single record in your organization's data, whereas a field like ``org_id`` will be + a constant across *all* your organization's data. +* Timestamp fields such as ``backend_timestamp`` are useful in ad-hoc queries, to look for data occurring before or + after a certain date, but are of limited usefulness over the span of time a watchlist may be running. +* A list of hashes (such as with ``process_sha256``) can be of limited value. They are inconvenient to keep current, + especially as software (whether legitimate or malicious) gets updated over time, but are definitely easier to manage + with ``equality`` IOCs. +* Counter fields (such as ``netconn_count``) can be useful with range queries to locate processes that are using a + large number of resources. For example, the query ``netconn_count:[500 TO *]`` will match only processes that make + a large number of network connections. +* When using ingress IOCs, be careful of errant characters in the ``values`` list, such as leading or trailing + whitespace or embedded newline characters. These errant characters may cause the IOCs to fail to match, leading to + false negative results. +* ``equality`` IOCs for IPv4 fields (e.g. ``netconn_remote_ipv4``) cannot support CIDR notation; full IP addresses + must be used. +* ``equality`` IOCs for IPv6 fields (e.g. ``netconn_remote_ipv6``) do not support standard or CIDR notation at this + time. All IPv6 addresses must omit colon characters, spell out all zeroes in the address, and represent all + alphabetic characters in uppercase. For example, "ff02::fb" becomes "FF0200000000000000000000000000FB". + +Feeds +----- +Another way of managing reports is to attach them to a *feed.* Feeds can contain multiple reports, and a feed can be +attached to a watchlist, effectively making the contents of the watchlist equivalent to the contents of the feed. + +Feeds are in effect “potentially-subscribable Watchlists”. A Feed has no effect on your organization until it is +subscribed to, by creating a Watchlist containing that feed. Once subscribed (and until it’s disabled or unsubscribed), +a watchlist will generate hits (and alerts if you have enabled them) for any matches against any of the IOCs in any of +that feed’s enabled reports. + +**Note:** The feeds that are created by these examples are *private feeds,* meaning they are only visible within an +organization and can be created by anyone with sufficient privileges in the organization. There are additional types +of feeds; *reserved feeds* can only be created by MSSPs, and *public feeds* can only be created or edited by +VMware Carbon Black. + +A new feed may be created as follows (assuming the new report for that feed is stored in the ``report`` variable): + +:: + + >>> from cbc_sdk.enterprise_edr import Feed + >>> builder = Feed.create(api, 'Suspicious Applications', 'http://example.com/location', + ... 'Any signs of suspicious applications running on our endpoints', 'external_threat_intel') + >>> builder.set_source_label('Where the info is coming from') + >>> builder.add_reports([report]) + >>> feed = builder.build() + >>> feed.save() + +If you have an existing feed, a new report may be added to it as follows (assuming the new report is stored in the +``report`` variable): + +:: + + >>> from cbc_sdk.enterprise_edr import Feed + >>> feed = cb.select(Feed, 'ABCDEFGHIJKLMNOPQRSTUVWX') + >>> feed.append_reports([report]) + +To update or delete an existing report in a feed, look for it in the feed's ``reports`` collection, then call the +``update()`` method on the report to replace its contents, or the ``delete()`` method on the report to delete it +entirely. The ``replace_reports()`` method on the ``Feed`` object may also be used, but caution must be taken, as +that method will replace *all* of the reports in a feed at once. + +To subscribe to a feed, a new watchlist must be created around it: + +:: + + >>> watchlist = Watchlist.create_from_feed(feed, "Subscribed feed", "Subscription to the new feed") + >>> watchlist.save() + +Limitations of Reports and Watchlists +------------------------------------- +Individual reports may contain no more than 10,000 IOCs. Reports containing more than 1,000 IOCs will not be editable +via the Carbon Black Cloud console UI, but may still be managed using APIs. + +Individual watchlists may contain no more than 10,000 reports. Any more than that may lead to timeouts when managing +the watchlist through the Carbon Black Cloud console UI, and possibly when managing it through APIs as well. diff --git a/examples/enterprise_edr/threat_intelligence/stix_taxii.py b/examples/enterprise_edr/threat_intelligence/stix_taxii.py index 6973c745a..4544042f9 100644 --- a/examples/enterprise_edr/threat_intelligence/stix_taxii.py +++ b/examples/enterprise_edr/threat_intelligence/stix_taxii.py @@ -193,8 +193,13 @@ def parse_collection_content(self, content_blocks, default_score=None): score = default_score else: score = self.config.default_score - for block in content_blocks: - yield from parse_stix(block.content, score) + + try: + for block in content_blocks: + yield from parse_stix(block.content, score) + except: + # Content Block failed or parsing issue continue with current progress + yield from () def import_collection(self, collection): """ diff --git a/requirements.txt b/requirements.txt index c3d55286b..85a1f5f3f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,6 +7,7 @@ prompt_toolkit pygments python-dateutil protobuf +schema solrq validators diff --git a/src/cbc_sdk/__init__.py b/src/cbc_sdk/__init__.py index 7e0c419a6..38c8f92e4 100644 --- a/src/cbc_sdk/__init__.py +++ b/src/cbc_sdk/__init__.py @@ -4,7 +4,7 @@ __author__ = 'Carbon Black Developer Network' __license__ = 'MIT' __copyright__ = 'Copyright 2020-2021 VMware Carbon Black' -__version__ = '1.3.1' +__version__ = '1.3.2' from .rest_api import CBCloudAPI from .cache import lru diff --git a/src/cbc_sdk/endpoint_standard/base.py b/src/cbc_sdk/endpoint_standard/base.py index 46b218dd5..9b92403d9 100644 --- a/src/cbc_sdk/endpoint_standard/base.py +++ b/src/cbc_sdk/endpoint_standard/base.py @@ -316,6 +316,10 @@ def _get_detailed_results(self): while still_fetching: result = self._cb.get_object(result_url, query_parameters=query_parameters) total_results = result.get('num_available', 0) + found_results = result.get('num_found', 0) + # if found is 0, then no enriched events + if found_results == 0: + return self if total_results != 0: results = result.get('results', []) self._info = results[0] diff --git a/src/cbc_sdk/enterprise_edr/models/feed.yaml b/src/cbc_sdk/enterprise_edr/models/feed.yaml index 27c0e4f42..ccb20baca 100644 --- a/src/cbc_sdk/enterprise_edr/models/feed.yaml +++ b/src/cbc_sdk/enterprise_edr/models/feed.yaml @@ -1,11 +1,9 @@ type: object required: - name - - owner - provider_url - summary - category - - access properties: name: type: string diff --git a/src/cbc_sdk/enterprise_edr/models/report.yaml b/src/cbc_sdk/enterprise_edr/models/report.yaml index 6ef3bad19..25875722a 100644 --- a/src/cbc_sdk/enterprise_edr/models/report.yaml +++ b/src/cbc_sdk/enterprise_edr/models/report.yaml @@ -1,6 +1,5 @@ type: object required: - - id - timestamp - title - description diff --git a/src/cbc_sdk/enterprise_edr/models/watchlist.yaml b/src/cbc_sdk/enterprise_edr/models/watchlist.yaml index 349909975..1df866c18 100644 --- a/src/cbc_sdk/enterprise_edr/models/watchlist.yaml +++ b/src/cbc_sdk/enterprise_edr/models/watchlist.yaml @@ -2,8 +2,6 @@ type: object required: - name - description - - create_timestamp - - last_update_timestamp properties: name: type: string diff --git a/src/cbc_sdk/enterprise_edr/threat_intelligence.py b/src/cbc_sdk/enterprise_edr/threat_intelligence.py index fae655f7b..5ca174134 100644 --- a/src/cbc_sdk/enterprise_edr/threat_intelligence.py +++ b/src/cbc_sdk/enterprise_edr/threat_intelligence.py @@ -14,12 +14,16 @@ """Model Classes for Enterprise Endpoint Detection and Response""" from __future__ import absolute_import + +import uuid + from cbc_sdk.errors import ApiError, InvalidObjectError, NonQueryableModel from cbc_sdk.base import CreatableModelMixin, MutableBaseModel, UnrefreshableModel, SimpleQuery import logging import time import validators +from schema import And, Optional, Schema, SchemaError log = logging.getLogger(__name__) @@ -29,7 +33,35 @@ class FeedModel(UnrefreshableModel, CreatableModelMixin, MutableBaseModel): """A common base class for models used by the Feed and Watchlist APIs.""" - pass + SCHEMA_IOCV2 = Schema( + { + "id": And(And(str, error="IOC field 'id' is not a string"), len), + "match_type": And(And(str, error="IOC field 'match_type' is not a string"), + And(lambda type: type in ["query", "equality", "regex"], + error="error in IOC 'match_type' value: Invalid match type")), + "values": And(And(list, error="IOC field 'values' is not a list"), + [And(str, error="IOC value is not a string")], len), + Optional("field"): And(str, error="IOC field 'field' is not a string"), + Optional("link"): And(str, error="IOC field 'link' is not a string") + } + ) + SCHEMA_REPORT = Schema( + { + "id": And(And(str, error="Report field 'id' is not a string"), len), + "timestamp": And(And(int, error="Report field 'timestamp' is not an integer"), + And(lambda n: n > 0, error="Timestamp cannot be negative")), + "title": And(And(str, error="Report field 'title' is not a string"), len), + "description": And(And(str, error="Report field 'description' is not a string"), len), + "severity": And(And(int, error="Report field 'severity' is not an integer"), + And(lambda n: 0 < n < 11, error="Severity value out of range")), + Optional("link"): And(str, error="Report field 'link' is not a string"), + Optional("tags"): And(And(list, error="Report field 'tags' is not a list"), + [And(str, error="Report tag is not a string")]), + "iocs_v2": And(And(list, error="Report field 'iocs_v2' is not a list"), [SCHEMA_IOCV2], + And(len, error="Report should have at least one IOC")), + Optional("visibility"): And(str, error="Report field 'visibility' is not a string") + } + ) class Watchlist(FeedModel): @@ -39,10 +71,6 @@ class Watchlist(FeedModel): urlobject_single = "/threathunter/watchlistmgr/v2/watchlist/{}" swagger_meta_file = "enterprise_edr/models/watchlist.yaml" - @classmethod - def _query_implementation(self, cb, **kwargs): - return WatchlistQuery(self, cb) - def __init__(self, cb, model_unique_id=None, initial_data=None): """ Initialize the Watchlist object. @@ -64,6 +92,164 @@ def __init__(self, cb, model_unique_id=None, initial_data=None): super(Watchlist, self).__init__(cb, model_unique_id=feed_id, initial_data=item, force_init=False, full_doc=True) + class WatchlistBuilder: + """Helper class allowing Watchlists to be assembled.""" + def __init__(self, cb, name): + """ + Creates a new WatchlistBuilder object. + + Args: + cb (CBCloudAPI): A reference to the CBCloudAPI object. + name (str): Name for the new watchlist. + """ + self._cb = cb + self._new_info = {"name": name, "tags_enabled": True, "alerts_enabled": False, "report_ids": []} + + def set_name(self, name): + """ + Sets the name for the new watchlist. + + Args: + name (str): New name for the watchlist. + + Returns: + WatchlistBuilder: This object. + """ + self._new_info['name'] = name + return self + + def set_description(self, description): + """ + Sets the description for the new watchlist. + + Args: + description (str): New description for the watchlist. + + Returns: + WatchlistBuilder: This object. + """ + self._new_info['description'] = description + return self + + def set_tags_enabled(self, flag): + """ + Sets whether tags will be enabled on the new watchlist. + + Args: + flag (bool): True to enable tags, False to disable them. Default is True. + + Returns: + WatchlistBuilder: This object. + """ + self._new_info['tags_enabled'] = bool(flag) + return self + + def set_alerts_enabled(self, flag): + """ + Sets whether alerts will be enabled on the new watchlist. + + Args: + flag (bool): True to enable alerts, False to disable them. Default is False. + + Returns: + WatchlistBuilder: This object. + """ + self._new_info['alerts_enabled'] = bool(flag) + return self + + def add_report_ids(self, report_ids): + """ + Adds report IDs to the watchlist. + + Args: + report_ids (list[str]): List of report IDs to add to the watchlist. + + Returns: + WatchlistBuilder: This object. + """ + self._new_info['report_ids'] += report_ids + return self + + def add_reports(self, reports): + """ + Adds reports to the watchlist. + + Args: + reports (list[Report]): List of reports to be added to the watchlist. + + Returns: + WatchlistBuilder: This object. + """ + id_values = [] + for report in reports: + if report._from_watchlist and 'id' in report._info: + report.validate() + id_values.append(report._info['id']) + return self.add_report_ids(id_values) + + def build(self): + """ + Builds the new Watchlist using information in the builder. The new watchlist must still be saved. + + Returns: + Watchlist: The new Watchlist. + """ + return Watchlist(self._cb, initial_data=self._new_info) + + @classmethod + def create(cls, cb, name): + """ + Starts creating a new Watchlist by returning a WatchlistBuilder that can be used to set attributes. + + Args: + cb (CBCloudAPI): A reference to the CBCloudAPI object. + name (str): Name for the new watchlist. + + Returns: + WatchlistBuilder: The builder for the new watchlist. Call build() to create the actual Watchlist. + """ + return Watchlist.WatchlistBuilder(cb, name) + + @classmethod + def create_from_feed(cls, feed, name=None, description=None, enable_alerts=False, enable_tags=True): + """ + Creates a new Watchlist that encapsulates a Feed. + + Args: + feed (Feed): The feed to be encapsulated by this Watchlist. + name (str): Name for the new watchlist. The default is to use the Feed name. + description (str): Description for the new watchlist. The default is to use the Feed summary. + enable_alerts (bool) - True to enable alerts, False to disable them. The default is False. + enable_tags (bool) - True to enable tags, False to disable them. The default is True. + + Returns: + Watchlist: A new Watchlist object, which must be saved to the server. + """ + return Watchlist(feed._cb, initial_data={ + "name": f"Feed {feed.name}" if not name else name, + "description": feed.summary if not description else description, + "tags_enabled": enable_tags, + "alerts_enabled": enable_alerts, + "classifier": { + "key": "feed_id", + "value": feed.id + } + }) + + @classmethod + def _query_implementation(self, cb, **kwargs): + """ + Returns the appropriate query object for Watchlists. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + **kwargs (dict): Not used, retained for compatibility. + + Returns: + WatchlistQuery: The query object for Watchlists. + """ + return WatchlistQuery(self, cb) + def save(self): """Saves this watchlist on the Enterprise EDR server. @@ -83,10 +269,11 @@ def save(self): return self def validate(self): - """Validates this watchlist's state. + """ + Checks to ensure this watchlist contains valid data. Raises: - InvalidObjectError: If the Watchlist's state is invalid. + InvalidObjectError: If the watchlist contains invalid data. """ super(Watchlist, self).validate() @@ -254,6 +441,30 @@ def reports(self): return reports_ + def add_report_ids(self, report_ids): + """ + Adds new report IDs to the watchlist. + + Args: + report_ids (list[str]): List of report IDs to be added to the watchlist. + """ + old_report_ids = self.report_ids if self._info.get('report_ids') else [] + self.update(report_ids=(old_report_ids + report_ids)) + + def add_reports(self, reports): + """ + Adds new reports to the watchlist. + + Args: + reports (list[Report]): List of reports to be added to the watchlist. + """ + report_ids = [] + for report in reports: + report.validate() + if report._from_watchlist: + report_ids.append(report._info['id']) + self.add_report_ids(report_ids) + class Feed(FeedModel): """Represents an Enterprise EDR feed's metadata.""" @@ -262,10 +473,6 @@ class Feed(FeedModel): primary_key = "id" swagger_meta_file = "enterprise_edr/models/feed.yaml" - @classmethod - def _query_implementation(self, cb, **kwargs): - return FeedQuery(self, cb) - def __init__(self, cb, model_unique_id=None, initial_data=None): """ Initialize the Feed object. @@ -300,8 +507,147 @@ def __init__(self, cb, model_unique_id=None, initial_data=None): self._reports = [Report(cb, initial_data=report, feed_id=feed_id) for report in reports] + class FeedBuilder: + """Helper class allowing Feeds to be assembled.""" + def __init__(self, cb, info): + """ + Creates a new FeedBuilder object. + + Args: + cb (CBCloudAPI): A reference to the CBCloudAPI object. + info (dict): The initial information for the new feed. + """ + self._cb = cb + self._new_feedinfo = info + self._reports = [] + + def set_name(self, name): + """ + Sets the name for the new feed. + + Args: + name (str): New name for the feed. + + Returns: + FeedBuilder: This object. + """ + self._new_feedinfo['name'] = name + return self + + def set_provider_url(self, provider_url): + """ + Sets the provider URL for the new feed. + + Args: + provider_url (str): New provider URL for the feed. + + Returns: + FeedBuilder: This object. + """ + self._new_feedinfo['provider_url'] = provider_url + return self + + def set_summary(self, summary): + """ + Sets the summary for the new feed. + + Args: + summary (str): New summary for the feed. + + Returns: + FeedBuilder: This object. + """ + self._new_feedinfo['summary'] = summary + return self + + def set_category(self, category): + """ + Sets the category for the new feed. + + Args: + category (str): New category for the feed. + + Returns: + FeedBuilder: This object. + """ + self._new_feedinfo['category'] = category + return self + + def set_source_label(self, source_label): + """ + Sets the source label for the new feed. + + Args: + source_label (str): New source label for the feed. + + Returns: + FeedBuilder: This object. + """ + self._new_feedinfo['source_label'] = source_label + return self + + def add_reports(self, reports): + """ + Adds new reports to the new feed. + + Args: + reports (list[Report]): New reports to be added to the feed. + + Returns: + FeedBuilder: This object. + """ + self._reports += reports + return self + + def build(self): + """ + Builds the new Feed. + + Returns: + Feed: The new Feed. + """ + report_data = [] + for report in self._reports: + report.validate() + report_data.append(report._info) + init_data = {'feedinfo': self._new_feedinfo, 'reports': report_data} + return Feed(self._cb, None, init_data) + + @classmethod + def create(cls, cb, name, provider_url, summary, category): + """ + Begins creating a new feed by making a FeedBuilder to hold the new feed data. + + Args: + cb (CBCloudAPI): A reference to the CBCloudAPI object. + name (str): Name for the new feed. + provider_url (str): Provider URL for the new feed. + summary (str): Summary for the new feed. + category (str): Category for the new feed. + + Returns: + FeedBuilder: The new FeedBuilder object to be used to create the feed. + """ + return Feed.FeedBuilder(cb, {'name': name, 'provider_url': provider_url, 'summary': summary, + 'category': category}) + + @classmethod + def _query_implementation(self, cb, **kwargs): + """ + Returns the appropriate query object for Feeds. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + **kwargs (dict): Not used, retained for compatibility. + + Returns: + FeedQuery: The query object for Feeds. + """ + return FeedQuery(self, cb) + def save(self, public=False): - """Saves this feed on the Enterprise EDR server. + """ + Saves this feed on the Enterprise EDR server. Arguments: public (bool): Whether to make the feed publicly available. @@ -311,9 +657,17 @@ def save(self, public=False): """ self.validate() + # Reports don't get assigned IDs by default when they get saved to a Feed. Make sure they have some. + report_data = [] + for report in self._reports: + info = report._info + if not info.get('id', None): + info['id'] = str(uuid.uuid4()) + report_data.append(info) + body = { 'feedinfo': self._info, - 'reports': [report._info for report in self._reports], + 'reports': report_data, } url = "/threathunter/feedmgr/v2/orgs/{}/feeds".format( @@ -324,17 +678,20 @@ def save(self, public=False): new_info = self._cb.post_object(url, body).json() self._info.update(new_info) + self._reports = [Report(self._cb, initial_data=report, feed_id=new_info['id']) + for report in report_data] return self def validate(self): - """Validates this feed's state. + """ + Checks to ensure this feed contains valid data. Raises: - InvalidObjectError: If the Feed's state is invalid. + InvalidObjectError: If the feed contains invalid data. """ super(Feed, self).validate() - if self.access not in ["public", "private"]: + if self.access and self.access not in ["public", "private"]: raise InvalidObjectError("access should be public or private") if not validators.url(self.provider_url): @@ -344,7 +701,8 @@ def validate(self): report.validate() def delete(self): - """Deletes this feed from the Enterprise EDR server. + """ + Deletes this feed from the Enterprise EDR server. Raises: InvalidObjectError: If `id` is missing. @@ -359,7 +717,8 @@ def delete(self): self._cb.delete_object(url) def update(self, **kwargs): - """Update this feed's metadata with the given arguments. + """ + Update this feed's metadata with the given arguments. Arguments: **kwargs (dict(str, str)): The fields to update. @@ -392,36 +751,56 @@ def update(self, **kwargs): @property def reports(self): - """Returns a list of Reports associated with this feed. + """ + Returns a list of Reports associated with this feed. Returns: Reports ([Report]): List of Reports in this Feed. """ - return self._cb.select(Report).where(feed_id=self.id) + if not self.id: + raise InvalidObjectError("missing feed ID") + self._reports = list(self._cb.select(Report).where(feed_id=self.id)) + return self._reports - def replace_reports(self, reports): - """Replace this Feed's Reports with the given Reports. + def _overwrite_reports(self, reports, raw_reports): + """ + Overwrites the Reports in this Feed with the given Reports. Arguments: reports ([Report]): List of Reports to replace existing Reports with. - - Raises: - InvalidObjectError: If `id` is missing. + raw_reports (list[dict]): List of raw report data to incorporate into the reports. """ - if not self.id: - raise InvalidObjectError("missing feed ID") - - rep_dicts = [report._info for report in reports] - body = {"reports": rep_dicts} + rep_dicts = [] + for report in reports: + report.validate() + rep_dicts.append(report._info) + body = {"reports": rep_dicts + raw_reports} url = "/threathunter/feedmgr/v2/orgs/{}/feeds/{}/reports".format( self._cb.credentials.org_key, self.id ) self._cb.post_object(url, body) + self._reports = reports + [Report(self._cb, initial_data=report, feed_id=self._info['id']) + for report in raw_reports] + + def replace_reports(self, reports): + """ + Replace this Feed's Reports with the given Reports. + + Arguments: + reports ([Report]): List of Reports to replace existing Reports with. + + Raises: + InvalidObjectError: If `id` is missing. + """ + if not self.id: + raise InvalidObjectError("missing feed ID") + self._overwrite_reports(reports, []) def append_reports(self, reports): - """Append the given Reports to this Feed's current Reports. + """ + Append the given Reports to this Feed's current Reports. Arguments: reports ([Report]): List of Reports to append to Feed. @@ -431,16 +810,60 @@ def append_reports(self, reports): """ if not self.id: raise InvalidObjectError("missing feed ID") + self._overwrite_reports(self._reports + reports, []) - rep_dicts = [report._info for report in reports] - rep_dicts += [report._info for report in self.reports] - body = {"reports": rep_dicts} + @classmethod + def _validate_report_rawdata(cls, report_data): + """ + Evaluate specified report raw data to make sure it's valid. - url = "/threathunter/feedmgr/v2/orgs/{}/feeds/{}/reports".format( - self._cb.credentials.org_key, - self.id - ) - self._cb.post_object(url, body) + Args: + report_data (list[dict]): List of raw report data specified as dicts. + + Raises: + InvalidObjectError: If validation fails for any part of the report data. + """ + for report in report_data: + try: + Feed.SCHEMA_REPORT.validate(report) + except SchemaError as e: + raise InvalidObjectError(e.errors[-1] if e.errors[-1] else e.autos[-1]) + + for ioc in report['iocs_v2']: + if ioc['match_type'] in ['equality', 'regex'] and 'field' not in ioc: + raise InvalidObjectError(f"IOC of type {ioc['match_type']} must have a 'field' value") + if ioc['match_type'] == 'query' and len(ioc['values']) != 1: + raise InvalidObjectError("query IOC should have one and only one value") + + def replace_reports_rawdata(self, report_data): + """ + Replace this Feed's Reports with the given reports, specified as raw data. + + Arguments: + report_data (list[dict]) A list of report data, formatted as per the API documentation for reports. + + Raises: + InvalidObjectError: If `id` is missing or validation of the data fails. + """ + if not self.id: + raise InvalidObjectError("missing feed ID") + Feed._validate_report_rawdata(report_data) + self._overwrite_reports([], report_data) + + def append_reports_rawdata(self, report_data): + """ + Append the given report data, formatted as per the API documentation for reports, to this Feed's Reports. + + Arguments: + report_data (list[dict]) A list of report data, formatted as per the API documentation for reports. + + Raises: + InvalidObjectError: If `id` is missing or validation of the data fails. + """ + if not self.id: + raise InvalidObjectError("missing feed ID") + Feed._validate_report_rawdata(report_data) + self._overwrite_reports(self._reports, report_data) class Report(FeedModel): @@ -451,6 +874,16 @@ class Report(FeedModel): @classmethod def _query_implementation(self, cb, **kwargs): + """ + Returns the appropriate query object for Reports. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + **kwargs (dict): Not used, retained for compatibility. + + Returns: + ReportQuery: The query object for Reports. + """ return ReportQuery(self, cb) def __init__(self, cb, model_unique_id=None, initial_data=None, @@ -465,7 +898,7 @@ def __init__(self, cb, model_unique_id=None, initial_data=None, feed_id (str): The ID of the feed this report is for. from_watchlist (str): The ID of the watchlist this report is for. """ - super(Report, self).__init__(cb, model_unique_id=initial_data.get("id"), + super(Report, self).__init__(cb, model_unique_id=initial_data.get("id", None), initial_data=initial_data, force_init=False, full_doc=True) @@ -482,14 +915,167 @@ def __init__(self, cb, model_unique_id=None, initial_data=None, self._iocs = IOC(cb, initial_data=self.iocs, report_id=self.id) if self.iocs_v2: self._iocs_v2 = [IOC_V2(cb, initial_data=ioc, report_id=self.id) for ioc in self.iocs_v2] + # this flag is set when we need to rebuild the 'ioc_v2' element of _info from the _iocs_v2 array + self._iocs_v2_need_rebuild = False + + class ReportBuilder: + """Helper class allowing Reports to be assembled.""" + def __init__(self, cb, report_body): + """ + Initialize a new ReportBuilder. + + Args: + cb (CBCloudAPI): A reference to the CBCloudAPI object. + report_body (dict): Partial report body which should be filled in with all "required" fields. + """ + self._cb = cb + self._report_body = report_body + self._iocs = [] + + def set_title(self, title): + """ + Set the title for the new report. + + Args: + title (str): New title for the report. + + Returns: + ReportBuilder: This object. + """ + self._report_body['title'] = title + return self + + def set_description(self, description): + """ + Set the description for the new report. + + Args: + description (str): New description for the report. + + Returns: + ReportBuilder: This object. + """ + self._report_body['description'] = description + return self + + def set_timestamp(self, timestamp): + """ + Set the timestamp for the new report. + + Args: + timestamp (int): New timestamp for the report. + + Returns: + ReportBuilder: This object. + """ + self._report_body['timestamp'] = timestamp + return self + + def set_severity(self, severity): + """ + Set the severity for the new report. + + Args: + severity (int): New severity for the report. + + Returns: + ReportBuilder: This object. + """ + self._report_body['severity'] = severity + return self + + def set_link(self, link): + """ + Set the link for the new report. + + Args: + link (str): New link for the report. + + Returns: + ReportBuilder: This object. + """ + self._report_body['link'] = link + return self + + def add_tag(self, tag): + """ + Adds a tag value to the new report. + + Args: + tag (str): The new tag for the object. + + Returns: + ReportBuilder: This object. + """ + self._report_body['tags'].append(tag) + return self + + def add_ioc(self, ioc): + """ + Adds an IOC to the new report. + + Args: + ioc (IOC_V2): The IOC to be added to the report. + + Returns: + ReportBuilder: This object. + """ + self._iocs.append(ioc) + return self + + def set_visibility(self, visibility): + """ + Set the visibility for the new report. + + Args: + visibility (str): New visibility for the report. + + Returns: + ReportBuilder: This object. + """ + self._report_body['visibility'] = visibility + return self + + def build(self): + """ + Builds the actual Report from the internal data of the ReportBuilder. + + Returns: + Report: The new Report. + """ + report = Report(self._cb, None, self._report_body) + report._iocs_v2 = self._iocs + report._iocs_v2_need_rebuild = True + return report + + @classmethod + def create(cls, cb, title, description, severity, timestamp=None, tags=None): + """ + Begin creating a new Report by returning a ReportBuilder. + + Args: + cb (CBCloudAPI): A reference to the CBCloudAPI object. + title (str): Title for the new report. + description (str): Description for the new report. + severity (int): Severity value for the new report. + timestamp (int): UNIX-epoch timestamp for the new report. If omitted, current time will be used. + tags (list[str]): Tags to be added to the report. If omitted, there will be none. + + Returns: + ReportBuilder: Reference to the ReportBuilder object. + """ + if not timestamp: + timestamp = int(time.time()) + return Report.ReportBuilder(cb, {'title': title, 'description': description, 'severity': severity, + 'timestamp': timestamp, 'tags': tags if tags else []}) def save_watchlist(self): - """Saves this report *as a watchlist report*. + """ + Saves this report *as a watchlist report*. Note: - This method **cannot** be used to save a feed report. To - save feed reports, create them with `cb.create` and use - `Feed.replace`. + This method **cannot** be used to save a feed report. To save feed reports, create them with `cb.create` + and use `Feed.replace`. Raises: InvalidObjectError: If Report.validate() fails. @@ -509,10 +1095,11 @@ def save_watchlist(self): return self def validate(self): - """Validates this report's state. + """ + Checks to ensure this report contains valid data. Raises: - InvalidObjectError: If the report's state is invalid + InvalidObjectError: If the report contains invalid data. """ super(Report, self).validate() @@ -521,6 +1108,9 @@ def validate(self): if self.iocs_v2: [ioc.validate() for ioc in self._iocs_v2] + if self._iocs_v2_need_rebuild: + self._info['iocs_v2'] = [ioc._info for ioc in self._iocs_v2] + self._iocs_v2_need_rebuild = False def update(self, **kwargs): """Update this Report with the given arguments. @@ -536,8 +1126,7 @@ def update(self, **kwargs): and this report is a Feed Report, or Report.validate() fails. Note: - The report's timestamp is always updated, regardless of whether - passed explicitly. + The report's timestamp is always updated, regardless of whether passed explicitly. >>> report.update(title="My new report title") """ @@ -564,8 +1153,9 @@ def update(self, **kwargs): if self.iocs: self._iocs = IOC(self._cb, initial_data=self.iocs, report_id=self.id) - if self.iocs_v2: + if self.iocs_v2 and 'iocs_v2' in kwargs: self._iocs_v2 = [IOC_V2(self._cb, initial_data=ioc, report_id=self.id) for ioc in self.iocs_v2] + self._iocs_v2_need_rebuild = False # NOTE(ww): Updating reports on the watchlist API appears to require # updated timestamps. @@ -755,6 +1345,42 @@ def iocs_(self): # methods. return self._iocs_v2 + def append_iocs(self, iocs): + """ + Append a list of IOCs to this Report. + + Args: + iocs (list[IOC_V2]): List of IOCs to be added. + """ + if self.iocs_v2: + self._iocs_v2 += iocs + else: + self._iocs_v2 = iocs + self._iocs_v2_need_rebuild = True + + def remove_iocs_by_id(self, ids_list): + """ + Remove IOCs from this report by specifying their IDs. + + Args: + ids_list (list[str]): List of IDs of the IOCs to be removed. + """ + if self.iocs_v2: + id_set = set(ids_list) + old_len = len(self._iocs_v2) + self._iocs_v2 = [ioc for ioc in self._iocs_v2 if ioc._info['id'] not in id_set] + self._iocs_v2_need_rebuild = (old_len > len(self._iocs_v2)) + + def remove_iocs(self, iocs): + """ + Remove a list of IOCs from this Report. + + Args: + iocs (list[IOC_V2]): List of IOCs to be removed. + """ + if self.iocs_v2: + self.remove_iocs_by_id([ioc._info['id'] for ioc in iocs]) + class ReportSeverity(FeedModel): """Represents severity information for a Watchlist Report.""" @@ -777,15 +1403,32 @@ def __init__(self, cb, initial_data=None): full_doc=True) def _query_implementation(self, cb, **kwargs): - raise NonQueryableModel("IOC does not support querying") + """ + Queries are not supported for report severity, so this raises an exception. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + **kwargs (dict): Additional arguments. + + Raises: + NonQueryableModel: Always. + """ + raise NonQueryableModel("ReportSeverity does not support querying") class IOC(FeedModel): - """Represents a collection of categorized IOCs.""" + """Represents a collection of categorized IOCs. These objects are officially deprecated and replaced by IOC_V2.""" swagger_meta_file = "enterprise_edr/models/iocs.yaml" def __init__(self, cb, model_unique_id=None, initial_data=None, report_id=None): - """Creates a new IOC instance. + """ + Creates a new IOC instance. + + Arguments: + cb (BaseAPI): Reference to API object used to communicate with the server. + model_unique_id (str): Unique ID of this IOC. + initial_data (dict): Initial data used to populate the IOC. + report_id (str): ID of the report this IOC belongs to (if this is a watchlist IOC). Raises: ApiError: If `initial_data` is None. @@ -799,13 +1442,24 @@ def __init__(self, cb, model_unique_id=None, initial_data=None, report_id=None): self._report_id = report_id def _query_implementation(self, cb, **kwargs): + """ + Queries are not supported for IOCs, so this raises an exception. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + **kwargs (dict): Additional arguments. + + Raises: + NonQueryableModel: Always. + """ raise NonQueryableModel("IOC does not support querying") def validate(self): - """Validates this IOC structure's state. + """ + Checks to ensure this IOC contains valid data. Raises: - InvalidObjectError: If the IOC structure's state is invalid. + InvalidObjectError: If the IOC contains invalid data. """ super(IOC, self).validate() @@ -832,7 +1486,14 @@ class IOC_V2(FeedModel): swagger_meta_file = "enterprise_edr/models/ioc_v2.yaml" def __init__(self, cb, model_unique_id=None, initial_data=None, report_id=None): - """Creates a new IOC_V2 instance. + """ + Creates a new IOC_V2 instance. + + Arguments: + cb (BaseAPI): Reference to API object used to communicate with the server. + model_unique_id (Any): Unused. + initial_data (dict): Initial data used to populate the IOC. + report_id (str): ID of the report this IOC belongs to (if this is a watchlist IOC). Raises: ApiError: If `initial_data` is None. @@ -847,13 +1508,141 @@ def __init__(self, cb, model_unique_id=None, initial_data=None, report_id=None): self._report_id = report_id def _query_implementation(self, cb, **kwargs): + """ + Queries are not supported for IOCs, so this raises an exception. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + **kwargs (dict): Additional arguments. + + Raises: + NonQueryableModel: Always. + """ raise NonQueryableModel("IOC_V2 does not support querying") + @classmethod + def create_query(cls, cb, iocid, query): + """ + Creates a new "query" IOC. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + iocid (str): ID for the new IOC. If this is None, a UUID will be generated for the IOC. + query (str): Query to be incorporated in this IOC. + + Returns: + IOC_V2: New IOC data structure. + + Raises: + ApiError: If the query string is not present. + """ + if not query: + raise ApiError("IOC must have a query string") + if not iocid: + iocid = str(uuid.uuid4()) + return IOC_V2(cb, iocid, {'id': iocid, 'match_type': 'query', 'values': [query]}) + + @classmethod + def create_equality(cls, cb, iocid, field, *values): + """ + Creates a new "equality" IOC. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + iocid (str): ID for the new IOC. If this is None, a UUID will be generated for the IOC. + field (str): Name of the field to be matched by this IOC. + *values (list(str)): String values to match against the value of the specified field. + + Returns: + IOC_V2: New IOC data structure. + + Raises: + ApiError: If there is not at least one value to match against. + """ + if not field: + raise ApiError('IOC must have a field name') + if len(values) == 0: + raise ApiError('IOC must have at least one value') + if not iocid: + iocid = str(uuid.uuid4()) + return IOC_V2(cb, iocid, {'id': iocid, 'match_type': 'equality', 'field': field, 'values': list(values)}) + + @classmethod + def create_regex(cls, cb, iocid, field, *values): + """ + Creates a new "regex" IOC. + + Args: + cb (BaseAPI): Reference to API object used to communicate with the server. + iocid (str): ID for the new IOC. If this is None, a UUID will be generated for the IOC. + field (str): Name of the field to be matched by this IOC. + *values (list(str)): Regular expression values to match against the value of the specified field. + + Returns: + IOC_V2: New IOC data structure. + + Raises: + ApiError: If there is not at least one regular expression to match against. + """ + if not field: + raise ApiError('IOC must have a field name') + if len(values) == 0: + raise ApiError('IOC must have at least one value') + if not iocid: + iocid = str(uuid.uuid4()) + return IOC_V2(cb, iocid, {'id': iocid, 'match_type': 'regex', 'field': field, 'values': list(values)}) + + @classmethod + def ipv6_equality_format(cls, input): + """ + Turns a canonically-formatted IPv6 address into a string suitable for use in an equality IOC. + + Args: + input (str): The IPv6 address to be translated. + + Returns: + str: The translated form of IPv6 address. + + Raises: + ApiError: If the string is not in valid format. + """ + def _check_components(array): + """If any component of an array is not valid for IPv6 (1-4 hex digits), raise an error""" + for element in array: + if len(element) not in range(1, 5): + raise ApiError('invalid address format') + for ch in element: + if ch not in '0123456789abcdefABCDEF': + raise ApiError('invalid address format') + + # try split on double colon first + segments = input.split('::', maxsplit=1) + if len(segments) == 2: + # take prefix and suffix part, add zeroes in between + prefix = segments[0].split(':') if segments[0] else [] + suffix = segments[1].split(':') if segments[1] else [] + if len(prefix) + len(suffix) >= 8: + raise ApiError('invalid address format') + _check_components(prefix) + _check_components(suffix) + num_blank = 8 - (len(prefix) + len(suffix)) + parts = prefix + (['0000'] * num_blank) + suffix + else: + # split all on single colon + parts = input.split(':') + if len(parts) != 8: + raise ApiError('invalid address format') + _check_components(parts) + # left pad all parts with zeroes + processed_parts = [('0000' + part)[-4:] for part in parts] + return "".join(processed_parts).upper() + def validate(self): - """Validates this IOC_V2's state. + """ + Checks to ensure this IOC contains valid data. Raises: - InvalidObjectError: If the IOC_V2's state is invalid. + InvalidObjectError: If the IOC contains invalid data. """ super(IOC_V2, self).validate() @@ -862,10 +1651,13 @@ def validate(self): @property def ignored(self): - """Returns whether or not this IOC is ignored + """ + Returns whether or not this IOC is ignored. + + Only watchlist IOCs have an ignore status. Returns: - (bool): True if the IOC is ignore, False otherwise. + bool: True if the IOC is ignored, False otherwise. Raises: InvalidObjectError: If this IOC is missing an `id` or is not a Watchlist IOC. @@ -889,7 +1681,8 @@ def ignored(self): return resp["ignored"] def ignore(self): - """Sets the ignore status on this IOC. + """ + Sets the ignore status on this IOC. Only watchlist IOCs have an ignore status. @@ -909,7 +1702,8 @@ def ignore(self): self._cb.put_object(url, None) def unignore(self): - """Removes the ignore status on this IOC. + """ + Removes the ignore status on this IOC. Only watchlist IOCs have an ignore status. diff --git a/src/cbc_sdk/enterprise_edr/ubs.py b/src/cbc_sdk/enterprise_edr/ubs.py index 58dc71a43..db1cc0876 100644 --- a/src/cbc_sdk/enterprise_edr/ubs.py +++ b/src/cbc_sdk/enterprise_edr/ubs.py @@ -108,7 +108,7 @@ def download_url(self, expiration_seconds=3600): else: return next((item.url for item in downloads.found - if self.sha256 == item.sha256), None) + if self.sha256.upper() == item.sha256.upper()), None) class Downloads(UnrefreshableModel): diff --git a/src/cbc_sdk/live_response_api.py b/src/cbc_sdk/live_response_api.py index 645141493..875fb74df 100644 --- a/src/cbc_sdk/live_response_api.py +++ b/src/cbc_sdk/live_response_api.py @@ -14,24 +14,21 @@ """The Live Response API and associated objects.""" from __future__ import absolute_import - +from collections import defaultdict +from concurrent.futures import _base, ThreadPoolExecutor import json +import logging +import queue import random +import shutil import string import threading import time -import logging -from collections import defaultdict - -import shutil from cbc_sdk.platform import Device from cbc_sdk.errors import TimeoutError, ObjectNotFoundError, ApiError -from concurrent.futures import _base from cbc_sdk import winerror -import queue - OS_LIVE_RESPONSE_ENUM = { "WINDOWS": 1, "LINUX": 2, @@ -93,7 +90,7 @@ class CbLRSessionBase(object): MAX_RETRY_COUNT = 5 - def __init__(self, cblr_manager, session_id, device_id, session_data=None): + def __init__(self, cblr_manager, session_id, device_id, session_data=None, thread_pool_count=5): """ Initialize the CbLRSessionBase. @@ -102,11 +99,14 @@ def __init__(self, cblr_manager, session_id, device_id, session_data=None): session_id (str): The ID of this session. device_id (int): The ID of the device (remote machine) we're connected to. session_data (dict): Additional session data. + thread_pool_count (int): number of workers for async commands (optional) """ self.session_id = session_id self.device_id = device_id self._cblr_manager = cblr_manager self._cb = cblr_manager._cb + self._async_executor = None + self._thread_pool_count = thread_pool_count # TODO: refcount should be in a different object in the scheduler self._refcount = 1 self._closed = False @@ -130,6 +130,51 @@ def __exit__(self, exc_type, exc_val, exc_tb): """ self.close() + def _async_submit(self, func, *args, **kwargs): + """ + Submit a task to the executor, creating it if it doesn't yet exist. + + Args: + func (func): A callable to be executed as a background task. + *args (list): Arguments to be passed to the callable. + **kwargs (dict): Keyword arguments to be passed to the callable. + + Returns: + Future: A future object representing the background task, which will pass along the result. + """ + if not self._async_executor: + self._async_executor = ThreadPoolExecutor(max_workers=self._thread_pool_count) + return self._async_executor.submit(func, args, kwargs) + + def command_status(self, command_id): + """ + Check the status of async command + + Args: + command_id (int): command_id + + Returns: + status of the command + """ + url = "{cblr_base}/sessions/{0}/commands/{1}".format(self.session_id, command_id, cblr_base=self.cblr_base) + res = self._cb.get_object(url) + return res["status"].upper() + + def cancel_command(self, command_id): + """ + Cancel command if it is in status PENDING. + + Args: + command_id (int): command_id + """ + url = "{cblr_base}/sessions/{0}/commands/{1}".format(self.session_id, command_id, cblr_base=self.cblr_base) + res = self._cb.get_object(url) + if res["status"].upper() == 'PENDING': + self._cb.delete_object(url) + else: + raise ApiError(f'Cannot cancel command in status {res["status"].upper()}.' + ' Only commands in status PENDING can be cancelled.') + def close(self): """Close the Live Response session.""" self._cblr_manager.close_session(self.device_id, self.session_id) @@ -138,7 +183,19 @@ def close(self): # # File operations # - def get_raw_file(self, file_name, timeout=None, delay=None): + def _submit_get_file(self, file_name): + """Helper function for submitting get file command""" + data = {"name": "get file", "path": file_name} + + resp = self._lr_post_command(data).json() + file_details = resp.get('file_details', None) + if file_details: + file_id = file_details.get('file_id', None) + command_id = resp.get('id', None) + return file_id, command_id + return None, None + + def get_raw_file(self, file_name, timeout=None, delay=None, async_mode=False): """ Retrieve contents of the specified file on the remote machine. @@ -146,25 +203,31 @@ def get_raw_file(self, file_name, timeout=None, delay=None): file_name (str): Name of the file to be retrieved. timeout (int): Timeout for the operation. delay (float): Delay in seconds to wait before command complete. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async + or object: Contains the data of the file. """ - data = {"name": "get file", "path": file_name} - - resp = self._lr_post_command(data).json() - file_details = resp.get('file_details', None) - if file_details: - file_id = file_details.get('file_id', None) - command_id = resp.get('id', None) - self._poll_command(command_id, timeout=timeout, delay=delay) + file_id, command_id = self._submit_get_file(file_name) + if file_id and command_id: + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._get_raw_file(command_id, + file_id, + timeout, + delay)) + else: + return self._get_raw_file(command_id, file_id, timeout, delay) - response = self._cb.session.get("{cblr_base}/sessions/{0}/files/{1}/content".format( - self.session_id, file_id, cblr_base=self.cblr_base), stream=True) - response.raw.decode_content = True - return response.raw + def _get_raw_file(self, command_id, file_id, timeout=None, delay=None): + self._poll_command(command_id, timeout=timeout, delay=delay) + response = self._cb.session.get("{cblr_base}/sessions/{0}/files/{1}/content".format( + self.session_id, file_id, cblr_base=self.cblr_base), stream=True) + response.raw.decode_content = True + return response.raw - def get_file(self, file_name, timeout=None, delay=None): + def get_file(self, file_name, timeout=None, delay=None, async_mode=False): """ Retrieve contents of the specified file on the remote machine. @@ -172,29 +235,43 @@ def get_file(self, file_name, timeout=None, delay=None): file_name (str): Name of the file to be retrieved. timeout (int): Timeout for the operation. delay (float): Delay in seconds to wait before command complete. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async str: Contents of the specified file. """ - fp = self.get_raw_file(file_name, timeout=timeout, delay=delay) - content = fp.read() - fp.close() + def _get_file(): + """Helper function to get the content of a file""" + fp = self._get_raw_file(command_id, file_id, timeout=timeout, delay=delay) + content = fp.read() + fp.close() + return content - return content + file_id, command_id = self._submit_get_file(file_name) + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: _get_file()) + return _get_file() - def delete_file(self, filename): + def delete_file(self, filename, async_mode=False): """ Delete the specified file name on the remote machine. Args: filename (str): Name of the file to be deleted. + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ data = {"name": "delete file", "path": filename} resp = self._lr_post_command(data).json() command_id = resp.get('id') + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id)) self._poll_command(command_id) - def put_file(self, infp, remote_filename): + def put_file(self, infp, remote_filename, async_mode=False): r""" Create a new file on the remote machine with the specified data. @@ -205,6 +282,10 @@ def put_file(self, infp, remote_filename): Args: infp (object): Python file-like containing data to upload to the remote endpoint. remote_filename (str): File name to create on the remote endpoint. + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ data = {"name": "put file", "path": remote_filename} file_id = self._upload_file(infp) @@ -212,9 +293,11 @@ def put_file(self, infp, remote_filename): resp = self._lr_post_command(data).json() command_id = resp.get('id') + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id)) self._poll_command(command_id) - def list_directory(self, dir_name): + def list_directory(self, dir_name, async_mode=False): r""" List the contents of a directory on the remote machine. @@ -242,26 +325,36 @@ def list_directory(self, dir_name): Args: dir_name (str): Directory to list. This parameter should end with the path separator. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async + or list: A list of dicts, each one describing a directory entry. - """ data = {"name": "directory list", "path": dir_name} resp = self._lr_post_command(data).json() command_id = resp.get("id") + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id).get("files", [])) return self._poll_command(command_id).get("files", []) - def create_directory(self, dir_name): + def create_directory(self, dir_name, async_mode=False): """ Create a directory on the remote machine. Args: dir_name (str): The new directory name. + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ data = {"name": "create directory", "path": dir_name} resp = self._lr_post_command(data).json() command_id = resp.get('id') + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id)) self._poll_command(command_id) def _pathsep(self): @@ -315,6 +408,7 @@ def walk(self, top, topdown=True, onerror=None, followlinks=False): r""" Perform a full directory walk with recursion into subdirectories on the remote machine. + Note: walk does not support async_mode due to its behaviour, it can only be invoked synchronously Example: >>> with c.select(Device, 1).lr_session() as lr_session: ... for entry in lr_session.walk(directory_name): @@ -363,20 +457,25 @@ def walk(self, top, topdown=True, onerror=None, followlinks=False): # Process operations # - def kill_process(self, pid): + def kill_process(self, pid, async_mode=False): """ Terminate a process on the remote machine. Args: pid (int): Process ID to be terminated. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async bool: True if success, False if failure. """ data = {"name": "kill", "pid": pid} resp = self._lr_post_command(data).json() command_id = resp.get('id') - + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id, + timeout=10, + delay=0.1)) try: self._poll_command(command_id, timeout=10, delay=0.1) except TimeoutError: @@ -385,7 +484,7 @@ def kill_process(self, pid): return True def create_process(self, command_string, wait_for_output=True, remote_output_file_name=None, - working_directory=None, wait_timeout=30, wait_for_completion=True): + working_directory=None, wait_timeout=30, wait_for_completion=True, async_mode=False): """ Create a new process on the remote machine with the specified command string. @@ -403,8 +502,10 @@ def create_process(self, command_string, wait_for_output=True, remote_output_fil working_directory (str): The working directory of the create process operation. wait_timeout (int): Timeout used for this command. wait_for_completion (bool): True to wait until the process is completed before returning. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async str: The output of the process. """ # process is: @@ -413,6 +514,19 @@ def create_process(self, command_string, wait_for_output=True, remote_output_fil # - wait for the process to complete # - get the temporary file from the endpoint # - delete the temporary file + def wait_to_complete_command(): + if wait_for_completion: + self._poll_command(command_id, timeout=wait_timeout) + + if wait_for_output: + # now the file is ready to be read + file_content = self.get_file(data["output_file"]) + # delete the file + self._lr_post_command({"name": "delete file", "path": data["output_file"]}) + + return file_content + else: + return None if wait_for_output: wait_for_completion = True @@ -431,22 +545,12 @@ def create_process(self, command_string, wait_for_output=True, remote_output_fil resp = self._lr_post_command(data).json() command_id = resp.get('id') - - if wait_for_completion: - self._poll_command(command_id, timeout=wait_timeout) - - if wait_for_output: - # now the file is ready to be read - - file_content = self.get_file(data["output_file"]) - # delete the file - self._lr_post_command({"name": "delete file", "path": data["output_file"]}) - - return file_content + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: wait_to_complete_command()) else: - return None + return wait_to_complete_command() - def list_processes(self): + def list_processes(self, async_mode=False): r""" List currently running processes on the remote machine. @@ -463,13 +567,20 @@ def list_processes(self): u'sid': u's-1-5-18', u'username': u'NT AUTHORITY\\SYSTEM'} + Args: + async_mode (bool): Flag showing whether the command should be executed asynchronously + Returns: + command_id, future if ran async + or list: A list of dicts describing the processes. """ data = {"name": "process list"} resp = self._lr_post_command(data).json() command_id = resp.get('id') - + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id).get( + "processes", [])) return self._poll_command(command_id).get("processes", []) # @@ -478,7 +589,7 @@ def list_processes(self): # returns dictionary with 2 entries ("values" and "sub_keys") # "values" is a list containing a dictionary for each registry value in the key # "sub_keys" is a list containing one entry for each sub_key - def list_registry_keys_and_values(self, regkey): + def list_registry_keys_and_values(self, regkey, async_mode=False): r""" Enumerate subkeys and values of the specified registry key on the remote machine. @@ -513,36 +624,49 @@ def list_registry_keys_and_values(self, regkey): Args: regkey (str): The registry key to enumerate. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async + or dict: A dictionary with two keys, 'sub_keys' (a list of subkey names) and 'values' (a list of dicts containing value data, name, and type). """ + def _list_registry_keys_and_values(): + """Helper function for list registry keys and values""" + raw_output = self._poll_command(command_id) + return {'values': raw_output.get('values', []), 'sub_keys': raw_output.get('sub_keys', [])} data = {"name": "reg enum key", "path": regkey} resp = self._lr_post_command(data).json() command_id = resp.get('id') - raw_output = self._poll_command(command_id) - return {'values': raw_output.get('values', []), 'sub_keys': raw_output.get('sub_keys', [])} + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: _list_registry_keys_and_values()) + return _list_registry_keys_and_values() # returns a list containing a dictionary for each registry value in the key - def list_registry_values(self, regkey): + def list_registry_values(self, regkey, async_mode=False): """ Enumerate all registry values from the specified registry key on the remote machine. Args: regkey (str): The registry key to enumerate. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async + or list: List of values for the registry key. """ data = {"name": "reg enum key", "path": regkey} resp = self._lr_post_command(data).json() command_id = resp.get('id') - + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id).get( + "values", [])) return self._poll_command(command_id).get("values", []) # returns a dictionary with the registry value - def get_registry_value(self, regkey): + def get_registry_value(self, regkey, async_mode=False): r""" Return the associated value of the specified registry key on the remote machine. @@ -553,17 +677,22 @@ def get_registry_value(self, regkey): Args: regkey (str): The registry key to retrieve. + async_mode (bool): Flag showing whether the command should be executed asynchronously Returns: + command_id, future if ran async + or dict: A dictionary with keys of: value_data, value_name, value_type. """ data = {"name": "reg query value", "path": regkey} resp = self._lr_post_command(data).json() command_id = resp.get('id') - + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id).get( + "value", {})) return self._poll_command(command_id).get("value", {}) - def set_registry_value(self, regkey, value, overwrite=True, value_type=None): + def set_registry_value(self, regkey, value, overwrite=True, value_type=None, async_mode=False): r""" Set a registry value on the specified registry key on the remote machine. @@ -576,6 +705,10 @@ def set_registry_value(self, regkey, value, overwrite=True, value_type=None): value (object): The value data. overwrite (bool): If True, any existing value will be overwritten. value_type (str): The type of value. Examples: REG_DWORD, REG_MULTI_SZ, REG_SZ + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ real_value = value if value_type is None: @@ -592,48 +725,68 @@ def set_registry_value(self, regkey, value, overwrite=True, value_type=None): "value_data": real_value} resp = self._lr_post_command(data).json() command_id = resp.get('id') + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id)) self._poll_command(command_id) - def create_registry_key(self, regkey): + def create_registry_key(self, regkey, async_mode=False): """ Create a new registry key on the remote machine. Args: regkey (str): The registry key to create. + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ data = {"name": "reg create key", "path": regkey} resp = self._lr_post_command(data).json() command_id = resp.get('id') + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id)) self._poll_command(command_id) - def delete_registry_key(self, regkey): + def delete_registry_key(self, regkey, async_mode=False): """ Delete a registry key on the remote machine. Args: regkey (str): The registry key to delete. + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ data = {"name": "reg delete key", "path": regkey} resp = self._lr_post_command(data).json() command_id = resp.get('id') + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id)) self._poll_command(command_id) - def delete_registry_value(self, regkey): + def delete_registry_value(self, regkey, async_mode=False): """ Delete a registry value on the remote machine. Args: regkey (str): The registry value to delete. + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ data = {"name": "reg delete value", "path": regkey} resp = self._lr_post_command(data).json() command_id = resp.get('id') + if async_mode: + return command_id, self._async_submit(lambda arg, kwarg: self._poll_command(command_id)) self._poll_command(command_id) # # Physical memory capture # - def memdump(self, local_filename, remote_filename=None, compress=False): + def memdump(self, local_filename, remote_filename=None, compress=False, async_mode=False): """ Perform a memory dump operation on the remote machine. @@ -641,11 +794,20 @@ def memdump(self, local_filename, remote_filename=None, compress=False): local_filename (str): Name of the file the memory dump will be transferred to on the local machine. remote_filename (str): Name of the file the memory dump will be stored in on the remote machine. compress (bool): True to compress the file on the remote system. + async_mode (bool): Flag showing whether the command should be executed asynchronously + + Returns: + command_id, future if ran async """ + def _memdump(): + """Helper function for memdump""" + dump_object.wait() + dump_object.get(local_filename) + dump_object.delete() dump_object = self.start_memdump(remote_filename=remote_filename, compress=compress) - dump_object.wait() - dump_object.get(local_filename) - dump_object.delete() + if async_mode: + return dump_object.memdump_id, self._async_submit(lambda arg, kwarg: _memdump()) + _memdump() def start_memdump(self, remote_filename=None, compress=True): """ @@ -1016,7 +1178,7 @@ class CbLRManagerBase(object): cblr_base = "" # override in subclass for each product cblr_session_cls = NotImplemented # override in subclass for each product - def __init__(self, cb, timeout=30, keepalive_sessions=False): + def __init__(self, cb, timeout=30, keepalive_sessions=False, thread_pool_count=5): """ Initialize the CbLRManagerBase object. @@ -1024,6 +1186,7 @@ def __init__(self, cb, timeout=30, keepalive_sessions=False): cb (BaseAPI): The CBC SDK object reference. timeout (int): Timeout to use for requests, in seconds. keepalive_sessions (bool): If True, "ping" sessions occasionally to ensure they stay alive. + thread_pool_count (int): number of workers for async commands (optional) """ self._timeout = timeout self._cb = cb @@ -1032,6 +1195,8 @@ def __init__(self, cb, timeout=30, keepalive_sessions=False): self._keepalive_sessions = keepalive_sessions self._init_poll_delay = 1 self._init_poll_timeout = 360 + self._async_executor = None + self._thread_pool_count = thread_pool_count if keepalive_sessions: self._cleanup_thread_running = True @@ -1042,6 +1207,22 @@ def __init__(self, cb, timeout=30, keepalive_sessions=False): self._job_scheduler = None + def _async_submit(self, func, *args, **kwargs): + """ + Submit a task to the executor, creating it if it doesn't yet exist. + + Args: + func (func): A callable to be executed as a background task. + *args (list): Arguments to be passed to the callable. + **kwargs (dict): Keyword arguments to be passed to the callable. + + Returns: + Future: A future object representing the background task, which will pass along the result. + """ + if not self._async_executor: + self._async_executor = ThreadPoolExecutor(max_workers=self._thread_pool_count) + return self._async_executor.submit(func, args, kwargs) + def submit_job(self, job, device): """ Submit a new job to be executed as a Live Response. @@ -1100,7 +1281,7 @@ def stop_keepalive_thread(self): self._cleanup_thread_running = False self._cleanup_thread_event.set() - def request_session(self, device_id): + def request_session(self, device_id, async_mode=False): """ Initiate a new Live Response session. @@ -1110,20 +1291,37 @@ def request_session(self, device_id): Returns: CbLRSessionBase: The new Live Response session. """ + def _return_existing_session(): + self._sessions[device_id]._refcount += 1 + return session + + def _get_session_obj(): + _, session_data = self._wait_create_session(device_id, session_id) + session = self.cblr_session_cls(self, session_id, device_id, session_data=session_data) + return session + + def _create_and_store_session(): + session = _get_session_obj() + self._sessions[device_id] = session + return session + if self._keepalive_sessions: with self._session_lock: if device_id in self._sessions: session = self._sessions[device_id] - self._sessions[device_id]._refcount += 1 + if async_mode: + return session.session_id, self._async_submit(lambda arg, kwarg: _return_existing_session()) + return _return_existing_session() else: - session_id, session_data = self._get_or_create_session(device_id) - session = self.cblr_session_cls(self, session_id, device_id, session_data=session_data) - self._sessions[device_id] = session + session_id = self._create_session(device_id) + if async_mode: + return session_id, self._async_submit(lambda arg, kwarg: _create_and_store_session()) + return _create_and_store_session() else: - session_id, session_data = self._get_or_create_session(device_id) - session = self.cblr_session_cls(self, session_id, device_id, session_data=session_data) - - return session + session_id = self._create_session(device_id) + if async_mode: + return session_id, self._async_submit(lambda arg, kwarg: _get_session_obj()) + return _get_session_obj() def close_session(self, device_id, session_id): """ @@ -1197,7 +1395,23 @@ def submit_job(self, job, device): def _get_or_create_session(self, device_id): session_id = self._create_session(device_id) + return self._wait_create_session(device_id, session_id) + + def session_status(self, session_id): + """ + Check the status of a lr session + + Args: + session_id (str): The id of the session. + + Returns: + str: Status of the session + """ + url = "{cblr_base}/sessions/{0}".format(session_id, cblr_base=self.cblr_base) + res = self._cb.get_object(url) + return res['status'].upper() + def _wait_create_session(self, device_id, session_id): try: res = poll_status(self._cb, "{cblr_base}/sessions/{0}".format(session_id, cblr_base=self.cblr_base), @@ -1222,8 +1436,7 @@ def _close_session(self, session_id): def _create_session(self, device_id): response = self._cb.post_object("{cblr_base}/sessions".format(cblr_base=self.cblr_base), {"device_id": device_id}).json() - session_id = response["id"] - return session_id + return response["id"] class GetFileJob(object): @@ -1285,6 +1498,8 @@ def poll_status(cb, url, desired_status="COMPLETE", timeout=None, delay=None): return res elif res["status"].upper() == "ERROR": raise LiveResponseError(res) + elif res["status"].upper() == "CANCELLED": + raise ApiError('The command has been cancelled.') else: time.sleep(delay) diff --git a/src/cbc_sdk/platform/devices.py b/src/cbc_sdk/platform/devices.py index 14a801627..e26c7f0e5 100644 --- a/src/cbc_sdk/platform/devices.py +++ b/src/cbc_sdk/platform/devices.py @@ -81,7 +81,7 @@ def _refresh(self): self._last_refresh_time = time.time() return True - def lr_session(self): + def lr_session(self, async_mode=False): """ Retrieve a Live Response session object for this Device. @@ -91,7 +91,7 @@ def lr_session(self): Raises: ApiError: If there is an error establishing a Live Response session for this Device. """ - return self._cb._request_lr_session(self._model_unique_id) + return self._cb._request_lr_session(self._model_unique_id, async_mode=async_mode) def background_scan(self, flag): """ diff --git a/src/cbc_sdk/platform/processes.py b/src/cbc_sdk/platform/processes.py index 6b7e2fb58..27fe64ab9 100644 --- a/src/cbc_sdk/platform/processes.py +++ b/src/cbc_sdk/platform/processes.py @@ -420,6 +420,10 @@ def _get_detailed_results(self): while still_fetching: result = self._cb.get_object(result_url, query_parameters=query_parameters) total_results = result.get('num_available', 0) + found_results = result.get('num_found', 0) + # if found is 0, then no details + if found_results == 0: + return self if total_results != 0: results = result.get('results', []) self._info = results[0] diff --git a/src/cbc_sdk/rest_api.py b/src/cbc_sdk/rest_api.py index 3911bbcf9..156bb53f7 100644 --- a/src/cbc_sdk/rest_api.py +++ b/src/cbc_sdk/rest_api.py @@ -101,8 +101,8 @@ def live_response(self): self._lr_scheduler = LiveResponseSessionManager(self) return self._lr_scheduler - def _request_lr_session(self, device_id): - return self.live_response.request_session(device_id) + def _request_lr_session(self, device_id, async_mode=False): + return self.live_response.request_session(device_id, async_mode=async_mode) # ---- Audit and Remediation diff --git a/src/tests/uat/enriched_events_uat.py b/src/tests/uat/enriched_events_uat.py index 28af1e712..fbc18be6a 100644 --- a/src/tests/uat/enriched_events_uat.py +++ b/src/tests/uat/enriched_events_uat.py @@ -194,7 +194,7 @@ def enriched_events_details(cb, event_id): print(f"event_id: {event_id}\n") events = cb.select(EnrichedEvent).where(event_id=event_id) - pprint(events[0].get_details()) + pprint(events[0].get_details()._info) print("\nCompare results manually with Postman") print("----------------------------------------------------------") diff --git a/src/tests/uat/live_response_api_async_uat.py b/src/tests/uat/live_response_api_async_uat.py new file mode 100644 index 000000000..05505b502 --- /dev/null +++ b/src/tests/uat/live_response_api_async_uat.py @@ -0,0 +1,389 @@ +#!/usr/bin/env python +# ******************************************************* +# Copyright (c) VMware, Inc. 2021. All Rights Reserved. +# SPDX-License-Identifier: MIT +# ******************************************************* +# * +# * DISCLAIMER. THIS PROGRAM IS PROVIDED TO YOU "AS IS" WITHOUT +# * WARRANTIES OR CONDITIONS OF ANY KIND, WHETHER ORAL OR WRITTEN, +# * EXPRESS OR IMPLIED. THE AUTHOR SPECIFICALLY DISCLAIMS ANY IMPLIED +# * WARRANTIES OR CONDITIONS OF MERCHANTABILITY, SATISFACTORY QUALITY, +# * NON-INFRINGEMENT AND FITNESS FOR A PARTICULAR PURPOSE. + +""" +To execute, a profile must be provided using the standard CBC Credentials. + +Processes: Live Response +* Create sesssion +* File commands + * Create directory + * Create/Put file in a directory + * List a directory + * Walk a directory + * Get a file + * Delete a file + * Delete a directory + * Memdump +* Process commands + * Create a process + * List processes + * Kill a process +* Registry commands + * List registry keys and values + * List registry values + * Create registry key + * Create registry value + * Delete registry value + * Delete registry key +* Close a session +""" + +# Standard library imports +import os +import os.path +from os import path +from pprint import pprint +import sys + +# Internal library imports +from cbc_sdk.helpers import build_cli_parser, get_cb_cloud_object +from cbc_sdk.platform import Device +from cbc_sdk.live_response_api import LiveResponseError, LiveResponseSessionManager +from cbc_sdk.errors import ApiError + +# CONSTANTS +HEADERS = {'X-Auth-Token': '', 'Content-Type': 'application/json'} +ORG_KEY = '' +HOSTNAME = '' +DEVICE_ID = 8612331 + +# Files constants +LOCAL_FILE = 'memdump.txt' +DIR = 'C:\\\\demo\\\\' +MEMDUMP_FILE = r"c:\memdump.txt" +FILE = r"c:\demo\test.txt" +FILE_CONTENT = 1000 * "This is just a test" + +# Registry constants +KEY_PATH = 'HKCU\\Environment\\Test' +KEY_PATH_PARENT = 'HKCU\\Environment' + +# Formatters +NEWLINES = 1 +DELIMITER = '-' +SYMBOLS = 48 + + +def flatten_response(response): + """Helper function to extract only the filenames from the list directory command""" + return [item['filename'] for item in response] + + +def flatten_processes_response(response): + """Helper function to extract only pids and paths from the list processes command""" + return [{item['process_pid']: item['process_path']} for item in response] + + +def flatten_processes_response_filter(response, process_name): + """Helper function to extract only pids and paths from the list processes command""" + return [item['process_pid'] for item in response if process_name in item['process_path']] + + +def asynchronous_version(cb): + """Asynchronous commands""" + print() + print(11 * ' ' + 'Async Live Response Commands') + print(SYMBOLS * DELIMITER) + + try: + device = cb.select(Device, DEVICE_ID) + except Exception: + # if this is not run on dev01, pick a random active device + devices = cb.select(Device).set_status(['ACTIVE']) + device = devices[0] + + manager = LiveResponseSessionManager(cb) + session_id, result = manager.request_session(device.id, async_mode=True) + print(f'Created Session {session_id}...............OK') + while True: + status = manager.session_status(session_id) + if status == 'ACTIVE': + break + print('Current status:', status) + + lr_session = result.result() + print(f'Active Session {lr_session.session_id}................OK') + print() + + print() + print(10 * ' ' + 'Live Response Files Commands') + print(SYMBOLS * DELIMITER) + + # those two command will not raise an exception, because they are async + _, result1 = lr_session.create_directory(DIR, async_mode=True) + _, result2 = lr_session.create_directory(DIR, async_mode=True) + # create directory + try: + # upon getting the result, one of the commands will raise a LiveResponseError exception + result1.result() + result2.result() + except LiveResponseError as ex: + if 'ERROR_ALREADY_EXISTS' in str(ex): + print('Command raised ERROR_ALREADY_EXISTS...........OK') + else: + raise + print('Create Dir....................................OK') + + # show that test.txt is not present in that directory + _, ldresult = lr_session.list_directory(DIR, async_mode=True) + + # upload the file on the server + _, uresult = lr_session.put_file(open("test.txt", "r"), FILE, async_mode=True) + uresult.result() + print('PUT File......................................OK') + + directories = flatten_response(ldresult.result()) + assert 'test.txt' not in directories, f'Dirs: {directories}' + + # show that test.txt is present in that directory + _, lresult = lr_session.list_directory(DIR, async_mode=True) + directories = flatten_response(lresult.result()) + assert 'test.txt' in directories, f'Dirs: {directories}' + + # start a few commands one after the other + _, gresult = lr_session.get_file(FILE, async_mode=True) + g2command_id, g2result = lr_session.get_file(FILE, async_mode=True) + _, dresult = lr_session.delete_file(FILE, async_mode=True) + lr_session.cancel_command(g2command_id) + print('Started async commands........................OK') + + gresult.result() + print('GET File......................................OK') + + exc_raise = False + try: + g2result.result() + except ApiError as ex: + assert 'The command has been cancelled.' in str(ex) + exc_raise = True + finally: + assert exc_raise + + # make sure the file is deleted before checking that it does not exist + dresult.result() + exc_raise = False + _, result = lr_session.get_file(FILE, async_mode=True) + try: + result.result() + except LiveResponseError as ex: + assert 'ERROR_FILE_NOT_FOUND' in str(ex) or 'ERROR_PATH_NOT_FOUND' in str(ex), f'Other error {ex}' + exc_raise = True + finally: + assert exc_raise + + print('DELETE File...................................OK') + # delete the directory too + lr_session.delete_file(DIR) + print('DELETE Dir....................................OK') + + print() + print(9 * ' ' + 'Live Response Process Commands') + print(SYMBOLS * DELIMITER) + + # list processes + _, lresult = lr_session.list_processes(async_mode=True) + # create infinite ping, that could be killed afterwards + _, cresult = lr_session.create_process(r'cmd.exe /c "ping.exe -t 127.0.0.1"', + wait_for_completion=False, + wait_for_output=False, + async_mode=True) + cresult.result() + processes = lresult.result() + + pprint(flatten_processes_response(processes)) + print('List Processes................................OK') + + _, lresult = lr_session.list_processes(async_mode=True) + processes = lresult.result() + found = False + for process in processes: + if 'ping.exe' in process['process_path']: + found = True + + # assert that indeed there is such a process + assert found + print('Create Process................................OK') + + # kill all of the processes that are for ping.exe + for pid in flatten_processes_response_filter(processes, 'ping.exe'): + print(f'Killing process with pid {pid}') + _, result = lr_session.kill_process(pid, async_mode=True) + result.result() + + print('Kill Process..................................OK') + + print() + print(6 * ' ' + 'Live Response Registry Keys Commands') + print(SYMBOLS * DELIMITER) + + _, kv_result = lr_session.list_registry_keys_and_values(KEY_PATH_PARENT, async_mode=True) + _, kv2_result = lr_session.list_registry_keys_and_values(KEY_PATH, async_mode=True) + kvccommand_id, kvc_result = lr_session.create_registry_key(KEY_PATH, async_mode=True) + + result = kv_result.result() + pprint(result) + print('List Registry Keys and Values.................OK') + + found = False + for item in result['values']: + if item['registry_name'] == 'Test': + found = True + break + # make sure the value for the key doesn't exist + assert found is False + + exists = True + try: + result = kv2_result.result() + print(result) + except LiveResponseError as ex: + # make sure such registry key does not exist + exists = not ('ERROR_FILE_NOT_FOUND' in str(ex)) + assert exists is False + print(f'Registry key with path {KEY_PATH} does not exist. Creating...') + + while True: + status = lr_session.command_status(kvccommand_id) + if status == 'COMPLETE': + break + + result = kvc_result.result() + + exists = True + try: + _, result = lr_session.list_registry_keys_and_values(KEY_PATH, async_mode=True) + result.result() + except LiveResponseError as ex: + exists = False + assert 'ERROR_FILE_NOT_FOUND' in str(ex) + # make sure that now the key exists + assert exists + print('Create Registry Key...........................OK') + + # create value for the key + _, result = lr_session.set_registry_value(KEY_PATH, 1, async_mode=True) + result.result() + _, fresult = lr_session.get_registry_value(KEY_PATH, async_mode=True) + result = fresult.result() + assert result['registry_data'] == '1' + assert result['registry_name'] == 'Test' + print('Create Registry Value.........................OK') + + pprint(result) + print('Get Registry Value............................OK') + + _, fresult = lr_session.list_registry_keys_and_values(KEY_PATH_PARENT, async_mode=True) + result = fresult.result() + pprint(result) + print('List Registry Keys and Values.................OK') + + found = False + for item in result['values']: + if item['registry_name'] == 'Test': + found = True + break + # make sure the value for the key exists + assert found is True + + _, fresult = lr_session.delete_registry_value(KEY_PATH, async_mode=True) + fresult.result() + + # this will not raise an exception, because the result is not obtained + _, dresult = lr_session.delete_registry_value(KEY_PATH, async_mode=True) + fresult.result() + deleted = False + try: + dresult.result() + except LiveResponseError as ex: + deleted = 'ERROR_FILE_NOT_FOUND' in str(ex) + + # make sure the value is deleted + assert deleted + print('Delete Registry Value.........................OK') + + _, fresult = lr_session.list_registry_keys_and_values(KEY_PATH_PARENT, async_mode=True) + result = fresult.result() + found = False + for item in result['values']: + if item['registry_name'] == 'Test': + found = True + break + # make sure the value for the key was deleted successfully + assert found is False + + _, fresult = lr_session.delete_registry_key(KEY_PATH, async_mode=True) + result = fresult.result() + _, fresult = lr_session.list_registry_keys_and_values(KEY_PATH, async_mode=True) + exists = True + try: + fresult.result() + except LiveResponseError as ex: + # make sure such registry key does not exist + exists = False + assert 'ERROR_FILE_NOT_FOUND' in str(ex) + assert exists is False, 'Registry key was not properly deleted.' + print('Delete Registry Key...........................OK') + + lr_session.close() + print(f'Deleted the session {lr_session.session_id}...........OK') + + +def setup(): + """Setup Function""" + # create the file that will be uploaded on the server, + # will be deleted once we are done with the test + fp = open("test.txt", "w") + fp.write(FILE_CONTENT) + fp.close() + + +def teardown(): + """Teardown Function""" + # clean-up actions after the tests + if path.exists(LOCAL_FILE): + os.remove(LOCAL_FILE) + if path.exists("test.txt"): + os.remove("test.txt") + + +def main(): + """Script entry point""" + global ORG_KEY + global HOSTNAME + parser = build_cli_parser() + args = parser.parse_args() + print_detail = args.verbose + + if print_detail: + print(f"args provided {args}") + + cb = get_cb_cloud_object(args) + HEADERS['X-Auth-Token'] = cb.credentials.token + ORG_KEY = cb.credentials.org_key + HOSTNAME = cb.credentials.url + + # setup actions + setup() + + # asynchronous version of the commands + asynchronous_version(cb) + + # cleanup actions + teardown() + + +if __name__ == "__main__": + try: + sys.exit(main()) + except KeyboardInterrupt: + print("\nInterrupted by user") diff --git a/src/tests/uat/live_response_api_uat.py b/src/tests/uat/live_response_api_uat.py index d66b2fadb..d0db02673 100644 --- a/src/tests/uat/live_response_api_uat.py +++ b/src/tests/uat/live_response_api_uat.py @@ -39,11 +39,11 @@ """ # Standard library imports -import sys import os import os.path from os import path from pprint import pprint +import sys # Internal library imports from cbc_sdk.helpers import build_cli_parser, get_cb_cloud_object @@ -89,23 +89,9 @@ def flatten_processes_response_filter(response, process_name): return [item['process_pid'] for item in response if process_name in item['process_path']] -def main(): - """Script entry point""" - global ORG_KEY - global HOSTNAME - parser = build_cli_parser() - args = parser.parse_args() - print_detail = args.verbose - - if print_detail: - print(f"args provided {args}") - - cb = get_cb_cloud_object(args) - HEADERS['X-Auth-Token'] = cb.credentials.token - ORG_KEY = cb.credentials.org_key - HOSTNAME = cb.credentials.url - - print(13 * ' ' + 'Live Response Commands') +def synchronous_version(cb): + """Synchronous commands""" + print(11 * ' ' + 'Sync Live Response Commands') print(SYMBOLS * DELIMITER) try: @@ -119,12 +105,6 @@ def main(): print(f'Created Session {lr_session.session_id}...............OK') print() - # create the file that will be uploaded on the server, - # will be deleted once we are done with the test - fp = open("test.txt", "w") - fp.write(FILE_CONTENT) - fp.close() - print() print(10 * ' ' + 'Live Response Files Commands') print(SYMBOLS * DELIMITER) @@ -309,14 +289,52 @@ def main(): assert exists is False, 'Registry key was not properly deleted.' print('Delete Registry Key...........................OK') + lr_session.close() + print(f'Deleted the session {lr_session.session_id}...........OK') + + +def setup(): + """Setup Function""" + # create the file that will be uploaded on the server, + # will be deleted once we are done with the test + fp = open("test.txt", "w") + fp.write(FILE_CONTENT) + fp.close() + + +def teardown(): + """Teardown Function""" # clean-up actions after the tests if path.exists(LOCAL_FILE): os.remove(LOCAL_FILE) if path.exists("test.txt"): os.remove("test.txt") - lr_session.close() - print(f'Deleted the session {lr_session.session_id}...........OK') + +def main(): + """Script entry point""" + global ORG_KEY + global HOSTNAME + parser = build_cli_parser() + args = parser.parse_args() + print_detail = args.verbose + + if print_detail: + print(f"args provided {args}") + + cb = get_cb_cloud_object(args) + HEADERS['X-Auth-Token'] = cb.credentials.token + ORG_KEY = cb.credentials.org_key + HOSTNAME = cb.credentials.url + + # setup actions + setup() + + # synchronous version of the commands + synchronous_version(cb) + + # cleanup actions + teardown() if __name__ == "__main__": diff --git a/src/tests/uat/platform_devices_uat.py b/src/tests/uat/platform_devices_uat.py index 028b90698..2b25254ea 100755 --- a/src/tests/uat/platform_devices_uat.py +++ b/src/tests/uat/platform_devices_uat.py @@ -53,6 +53,14 @@ def search_devices_api_call(): return requests.post(search_url, json=data, headers=HEADERS) +def normalize_results(result): + """Helper function to sort dictionary values""" + for item in result: + if item.get('device_meta_data_item_list'): + dl = item['device_meta_data_item_list'] + item['device_meta_data_item_list'] = sorted(dl, key=lambda i: (i.get('key_name'))) + + def search_devices(cb): """Verify that the SDK returns the same result for Search Devices as the respective API call""" api_search = search_devices_api_call().json() @@ -69,6 +77,8 @@ def search_devices(cb): 'Test Failed: SDK call returns different number of devices. ' \ 'Expected: {}, Actual: {}'.format(api_search['num_found'], query._total_results) + normalize_results(api_search['results']) + normalize_results(sdk_results) assert sdk_results == api_search['results'], 'Test Failed: SDK call returns different data. '\ 'Expected: {}, Actual: {}'.format(api_search['results'], sdk_results) diff --git a/src/tests/unit/endpoint_standard/test_endpoint_standard_enriched_events.py b/src/tests/unit/endpoint_standard/test_endpoint_standard_enriched_events.py index 0bf4c34cf..2d1a4cfe4 100644 --- a/src/tests/unit/endpoint_standard/test_endpoint_standard_enriched_events.py +++ b/src/tests/unit/endpoint_standard/test_endpoint_standard_enriched_events.py @@ -8,6 +8,7 @@ from cbc_sdk.errors import ApiError from tests.unit.fixtures.CBCSDKMock import CBCSDKMock from tests.unit.fixtures.endpoint_standard.mock_enriched_events import ( + GET_ENRICHED_EVENTS_SEARCH_JOB_RESULTS_ZERO, POST_ENRICHED_EVENTS_SEARCH_JOB_RESP, GET_ENRICHED_EVENTS_SEARCH_JOB_RESULTS_RESP_1, GET_ENRICHED_EVENTS_SEARCH_JOB_RESULTS_RESP_2, @@ -156,6 +157,33 @@ def test_enriched_event_select_details_sync(cbcsdk_mock): assert results.process_pid[0] == 2000 +def test_enriched_event_select_details_sync_zero(cbcsdk_mock): + """Testing EnrichedEvent Querying with get_details""" + cbcsdk_mock.mock_request("POST", "/api/investigate/v2/orgs/test/enriched_events/search_job", + POST_ENRICHED_EVENTS_SEARCH_JOB_RESP) + cbcsdk_mock.mock_request("GET", + "/api/investigate/v1/orgs/test/enriched_events/search_jobs/08ffa932-b633-4107-ba56-8741e929e48b", # noqa: E501 + GET_ENRICHED_EVENTS_SEARCH_JOB_RESULTS_RESP) + cbcsdk_mock.mock_request("GET", + "/api/investigate/v2/orgs/test/enriched_events/search_jobs/08ffa932-b633-4107-ba56-8741e929e48b/results", # noqa: E501 + GET_ENRICHED_EVENTS_SEARCH_JOB_RESULTS_RESP_1) + cbcsdk_mock.mock_request("POST", "/api/investigate/v2/orgs/test/enriched_events/detail_jobs", + POST_ENRICHED_EVENTS_SEARCH_JOB_RESP) + cbcsdk_mock.mock_request("GET", + "/api/investigate/v2/orgs/test/enriched_events/detail_jobs/08ffa932-b633-4107-ba56-8741e929e48b", # noqa: E501 + GET_ENRICHED_EVENTS_SEARCH_JOB_RESULTS_RESP) + cbcsdk_mock.mock_request("GET", + "/api/investigate/v2/orgs/test/enriched_events/detail_jobs/08ffa932-b633-4107-ba56-8741e929e48b/results", # noqa: E501 + GET_ENRICHED_EVENTS_SEARCH_JOB_RESULTS_ZERO) + + s_api = cbcsdk_mock.api + events = s_api.select(EnrichedEvent).where(process_pid=2000) + event = events[0] + results = event.get_details() + assert results['device_name'] is not None + assert results.get('alert_id') is None + + def test_enriched_event_select_compound(cbcsdk_mock): """Testing EnrichedEvent Querying with select() and more complex criteria""" cbcsdk_mock.mock_request("POST", "/api/investigate/v2/orgs/test/enriched_events/search_job", diff --git a/src/tests/unit/enterprise_edr/test_enterprise_edr_threatintel.py b/src/tests/unit/enterprise_edr/test_enterprise_edr_threatintel.py index d6c8e240f..4dbf3532f 100644 --- a/src/tests/unit/enterprise_edr/test_enterprise_edr_threatintel.py +++ b/src/tests/unit/enterprise_edr/test_enterprise_edr_threatintel.py @@ -1,8 +1,11 @@ """Testing Watchlist, Report, Feed objects of cbc_sdk.enterprise_edr""" +import copy import pytest import logging -from cbc_sdk.enterprise_edr import Watchlist, Report, Feed +import re +from contextlib import ExitStack as does_not_raise +from cbc_sdk.enterprise_edr import Watchlist, Report, Feed, IOC_V2 from cbc_sdk.errors import InvalidObjectError, ApiError from cbc_sdk.rest_api import CBCloudAPI from tests.unit.fixtures.CBCSDKMock import CBCSDKMock @@ -15,9 +18,27 @@ REPORT_GET_RESP, FEED_GET_RESP, FEED_GET_SPECIFIC_RESP, - FEED_GET_SPECIFIC_FROM_WATCHLIST_RESP) + FEED_GET_SPECIFIC_FROM_WATCHLIST_RESP, + IOC_GET_IGNORED, + REPORT_BUILT_VIA_BUILDER, + REPORT_INIT, + REPORT_GET_IGNORED, + REPORT_GET_SEVERITY, + REPORT_UPDATE_AFTER_ADD_IOC, + REPORT_UPDATE_AFTER_REMOVE_IOC, + FEED_BUILT_VIA_BUILDER, + FEED_INIT, + FEED_UPDATE_INFO_1, + REPORT_INIT_2, + WATCHLIST_BUILDER_IN, + WATCHLIST_BUILDER_OUT, + WATCHLIST_FROM_FEED_IN, + WATCHLIST_FROM_FEED_OUT, + ADD_REPORT_IDS_LIST, + ADD_REPORTS_LIST) log = logging.basicConfig(format='%(asctime)s %(levelname)s:%(message)s', level=logging.DEBUG, filename='log.txt') +GUID_PATTERN = '[0-9A-Fa-f]{8}(-[0-9A-Fa-f]{4}){3}-[0-9A-Fa-f]{12}' @pytest.fixture(scope="function") @@ -330,3 +351,673 @@ def test_feed_query_specific(cbcsdk_mock): feed = api.select(Feed, "pv65TYVQy8YWMX9KsQUg") assert isinstance(feed, Feed) assert feed.id == "pv65TYVQy8YWMX9KsQUg" + + +def test_create_query_ioc(cb): + """Tests the creation of a 'query' IOC.""" + with pytest.raises(ApiError): + IOC_V2.create_query(cb, "foo", None) + + ioc = IOC_V2.create_query(cb, "foo", "process_name:evil.exe") + assert ioc._info == {'id': 'foo', 'match_type': 'query', 'values': ['process_name:evil.exe']} + ioc = IOC_V2.create_query(cb, None, "process_name:evil.exe") + assert ioc.match_type == 'query' + assert ioc.values == ['process_name:evil.exe'] + assert re.fullmatch(GUID_PATTERN, ioc.id) + + +def test_create_equality_ioc(cb): + """Tests the creation of an 'equality' IOC.""" + with pytest.raises(ApiError): + IOC_V2.create_equality(cb, "foo", None, "Alpha") + with pytest.raises(ApiError): + IOC_V2.create_equality(cb, "foo", "process_name") + + ioc = IOC_V2.create_equality(cb, "foo", "process_name", "Alpha", "Bravo", "Charlie") + assert ioc._info == {'id': 'foo', 'match_type': 'equality', 'field': 'process_name', + 'values': ['Alpha', 'Bravo', 'Charlie']} + ioc = IOC_V2.create_equality(cb, None, "process_name", "Alpha") + assert ioc.match_type == 'equality' + assert ioc.field == 'process_name' + assert ioc.values == ['Alpha'] + assert re.fullmatch(GUID_PATTERN, ioc.id) + + +def test_create_regex_ioc(cb): + """Tests the creation of a 'regex' IOC.""" + with pytest.raises(ApiError): + IOC_V2.create_regex(cb, "foo", None, "Alpha") + with pytest.raises(ApiError): + IOC_V2.create_regex(cb, "foo", "process_name") + + ioc = IOC_V2.create_regex(cb, "foo", "process_name", "Alpha", "Bravo", "Charlie") + assert ioc._info == {'id': 'foo', 'match_type': 'regex', 'field': 'process_name', + 'values': ['Alpha', 'Bravo', 'Charlie']} + ioc = IOC_V2.create_regex(cb, None, "process_name", "Alpha") + assert ioc.match_type == 'regex' + assert ioc.field == 'process_name' + assert ioc.values == ['Alpha'] + assert re.fullmatch(GUID_PATTERN, ioc._info['id']) + + +def test_ioc_read_ignored(cbcsdk_mock): + """Tests reading the ignore status of an IOC.""" + cbcsdk_mock.mock_request("GET", "/threathunter/watchlistmgr/v3/orgs/test/reports/a1b2/iocs/foo/ignore", + IOC_GET_IGNORED) + api = cbcsdk_mock.api + ioc = IOC_V2.create_equality(api, "foo", "process_name", "Alpha") + with pytest.raises(InvalidObjectError): + assert ioc.ignored + ioc._report_id = "a1b2" + assert ioc.ignored + ioc._info['id'] = None + with pytest.raises(InvalidObjectError): + assert ioc.ignored + + +def test_ioc_set_ignored(cbcsdk_mock): + """Tests setting the ignore status of an IOC.""" + cbcsdk_mock.mock_request("PUT", "/threathunter/watchlistmgr/v3/orgs/test/reports/a1b2/iocs/foo/ignore", + IOC_GET_IGNORED) + api = cbcsdk_mock.api + ioc = IOC_V2.create_equality(api, "foo", "process_name", "Alpha") + with pytest.raises(InvalidObjectError): + ioc.ignore() + ioc._report_id = "a1b2" + ioc.ignore() + ioc._info['id'] = None + with pytest.raises(InvalidObjectError): + ioc.ignore() + + +def test_ioc_clear_ignored(cbcsdk_mock): + """Tests clearing the ignore status of an IOC.""" + cbcsdk_mock.mock_request("DELETE", "/threathunter/watchlistmgr/v3/orgs/test/reports/a1b2/iocs/foo/ignore", + CBCSDKMock.StubResponse(None, 204)) + api = cbcsdk_mock.api + ioc = IOC_V2.create_equality(api, "foo", "process_name", "Alpha") + with pytest.raises(InvalidObjectError): + ioc.unignore() + ioc._report_id = "a1b2" + ioc.unignore() + ioc._info['id'] = None + with pytest.raises(InvalidObjectError): + ioc.unignore() + + +@pytest.mark.parametrize("input, output, expectation", [ + ('::1', '00000000000000000000000000000001', does_not_raise()), + ('2601:280:c300:294::6bd', '26010280C300029400000000000006BD', does_not_raise()), + ('fd9f:6dd0:e01b:0:1b7:daf8:b54d:4916', 'FD9F6DD0E01B000001B7DAF8B54D4916', does_not_raise()), + ('fe80::3c0c:73f:845b:8243', 'FE800000000000003C0C073F845B8243', does_not_raise()), + ('fd9f:6dd0:e01b:0:1b7:daf8:4916', None, pytest.raises(ApiError)), + ('fd9f:6dd0:e01b:0:1b7:daf8:b54d:4916:909', None, pytest.raises(ApiError)), + ('fd9f:6dd0:e01b:0:1b7:daef8:b54d:4916', None, pytest.raises(ApiError)), + ('fe80:3:C0:8080::3c0c:73f:845b:8243', None, pytest.raises(ApiError)), + ('2601::c300:294::6bd', None, pytest.raises(ApiError)), + ('fe80::3c0c:73g:845b:8243', None, pytest.raises(ApiError)), + ("10.8.16.254", None, pytest.raises(ApiError)) +]) +def test_ipv6_equality_format(input, output, expectation): + """Tests the ipv6_equality_format() function.""" + with expectation: + assert IOC_V2.ipv6_equality_format(input) == output + + +def test_report_builder_save_watchlist(cbcsdk_mock): + """Tests the operation of a ReportBuilder and saving the report as a watchlist report.""" + my_info = None + + def on_post(url, body, **kwargs): + assert my_info + assert body == my_info + my_info['id'] = "AaBbCcDdEeFfGg" + return my_info + + cbcsdk_mock.mock_request("POST", "/threathunter/watchlistmgr/v3/orgs/test/reports", on_post) + api = cbcsdk_mock.api + builder = Report.create(api, "NotReal", "Not real description", 2) + builder.set_title("ReportTitle").set_description("The report description").set_timestamp(1234567890) + builder.set_severity(5).set_link('https://example.com').add_tag("Alpha").add_tag("Bravo") + builder.add_ioc(IOC_V2.create_equality(api, "foo", "process_name", "evil.exe")) + builder.add_ioc(IOC_V2.create_equality(api, "bar", "netconn_ipv4", "10.29.99.1")) + builder.set_visibility("visible") + report = builder.build() + report.validate() + my_info = copy.deepcopy(report._info) + assert my_info == REPORT_BUILT_VIA_BUILDER + report.save_watchlist() + assert report._from_watchlist + assert report._info['id'] == "AaBbCcDdEeFfGg" + + +@pytest.mark.parametrize("init_data, feed, watchlist, do_request, expectation, result", [ + (REPORT_INIT, None, True, True, does_not_raise(), True), + (REPORT_INIT, None, False, False, pytest.raises(InvalidObjectError), True), + (REPORT_BUILT_VIA_BUILDER, None, True, False, pytest.raises(InvalidObjectError), True) +]) +def test_report_get_ignored(cbcsdk_mock, init_data, feed, watchlist, do_request, expectation, result): + """Tests the operation of the report.ignored() method.""" + if do_request: + cbcsdk_mock.mock_request("GET", "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7/ignore", REPORT_GET_IGNORED) + api = cbcsdk_mock.api + report = Report(api, None, init_data, feed, watchlist) + with expectation: + assert report.ignored == result + + +@pytest.mark.parametrize("init_data, feed, watchlist, do_request, expectation", [ + (REPORT_INIT, None, True, True, does_not_raise()), + (REPORT_INIT, None, False, False, pytest.raises(InvalidObjectError)), + (REPORT_BUILT_VIA_BUILDER, None, True, False, pytest.raises(InvalidObjectError)) +]) +def test_report_set_ignored(cbcsdk_mock, init_data, feed, watchlist, do_request, expectation): + """Tests the operation of the report.ignore() method.""" + if do_request: + cbcsdk_mock.mock_request("PUT", "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7/ignore", REPORT_GET_IGNORED) + api = cbcsdk_mock.api + report = Report(api, None, init_data, feed, watchlist) + with expectation: + report.ignore() + + +@pytest.mark.parametrize("init_data, feed, watchlist, do_request, expectation", [ + (REPORT_INIT, None, True, True, does_not_raise()), + (REPORT_INIT, None, False, False, pytest.raises(InvalidObjectError)), + (REPORT_BUILT_VIA_BUILDER, None, True, False, pytest.raises(InvalidObjectError)) +]) +def test_report_clear_ignored(cbcsdk_mock, init_data, feed, watchlist, do_request, expectation): + """Tests the operation of the report.unignore() method.""" + if do_request: + cbcsdk_mock.mock_request("DELETE", "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7/ignore", + CBCSDKMock.StubResponse(None, 204)) + api = cbcsdk_mock.api + report = Report(api, None, init_data, feed, watchlist) + with expectation: + report.unignore() + + +@pytest.mark.parametrize("init_data, feed, watchlist, call_url, expectation", [ + (REPORT_INIT, None, True, "/threathunter/watchlistmgr/v3/orgs/test/reports/69e2a8d0-bc36-4970-9834-8687efe1aff7", + does_not_raise()), + (REPORT_INIT, "qwertyuiop", False, "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7", does_not_raise()), + (REPORT_INIT, None, False, None, pytest.raises(InvalidObjectError)), + (REPORT_BUILT_VIA_BUILDER, None, True, None, pytest.raises(InvalidObjectError)), + (REPORT_BUILT_VIA_BUILDER, "qwertyuiop", False, None, pytest.raises(InvalidObjectError)) +]) +def test_report_delete(cbcsdk_mock, init_data, feed, watchlist, call_url, expectation): + """Tests the operation of the report.delete() method.""" + if call_url: + cbcsdk_mock.mock_request("DELETE", call_url, CBCSDKMock.StubResponse(None, 204)) + api = cbcsdk_mock.api + report = Report(api, None, init_data, feed, watchlist) + with expectation: + report.delete() + + +@pytest.mark.parametrize("init_data, feed, watchlist, do_request, expectation, result", [ + (REPORT_INIT, "qwertyuiop", False, True, does_not_raise(), 8), + (REPORT_INIT, None, True, False, pytest.raises(InvalidObjectError), -1), + (REPORT_BUILT_VIA_BUILDER, "qwertyuiop", False, False, pytest.raises(InvalidObjectError), -1) +]) +def test_report_get_custom_severity(cbcsdk_mock, init_data, feed, watchlist, do_request, expectation, result): + """Tests getting the custom severity for a report.""" + if do_request: + cbcsdk_mock.mock_request("GET", "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7/severity", REPORT_GET_SEVERITY) + api = cbcsdk_mock.api + report = Report(api, None, init_data, feed, watchlist) + with expectation: + severity = report.custom_severity + assert severity.report_id == report.id + assert severity.severity == result + + +@pytest.mark.parametrize("init_data, feed, watchlist, input, http_method, http_return, expectation", [ + (REPORT_INIT, "qwertyuiop", False, 8, "PUT", REPORT_GET_SEVERITY, does_not_raise()), + (REPORT_INIT, "qwertyuiop", False, None, "DELETE", CBCSDKMock.StubResponse(None, 204), does_not_raise()), + (REPORT_INIT, None, True, 8, None, None, pytest.raises(InvalidObjectError)), + (REPORT_BUILT_VIA_BUILDER, "qwertyuiop", False, 8, None, None, pytest.raises(InvalidObjectError)) +]) +def test_report_set_custom_severity(cbcsdk_mock, init_data, feed, watchlist, input, http_method, + http_return, expectation): + """Tests setting the custom severity for a report.""" + if http_method: + cbcsdk_mock.mock_request(http_method, "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7/severity", http_return) + api = cbcsdk_mock.api + report = Report(api, None, init_data, feed, watchlist) + with expectation: + report.custom_severity = input + + +def test_report_add_ioc(cbcsdk_mock): + """Test appending a new IOC and then updating the report.""" + def on_put(url, body, **kwargs): + match_data = copy.deepcopy(REPORT_UPDATE_AFTER_ADD_IOC) + match_data['timestamp'] = body['timestamp'] + assert body == match_data + return body + + cbcsdk_mock.mock_request("PUT", "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7", on_put) + api = cbcsdk_mock.api + report = Report(api, None, copy.deepcopy(REPORT_INIT), None, True) + report.append_iocs([IOC_V2.create_query(api, "quux", "filemod_name: \"audio.dat\"")]) + report.update() + assert len(report.iocs_) == 3 + + +def test_report_remove_ioc(cbcsdk_mock): + """Test removing an IOC and then updating the report.""" + def on_put(url, body, **kwargs): + match_data = copy.deepcopy(REPORT_UPDATE_AFTER_REMOVE_IOC) + match_data['timestamp'] = body['timestamp'] + assert body == match_data + return body + + cbcsdk_mock.mock_request("PUT", "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7", on_put) + api = cbcsdk_mock.api + report = Report(api, None, copy.deepcopy(REPORT_INIT), None, True) + ioc = report.iocs_[0] + assert ioc.id == "foo" + report.remove_iocs([ioc]) + report.update() + assert len(report.iocs_) == 1 + + +def test_report_remove_nonexistent_ioc_id(cbcsdk_mock): + """Test removing an IOC by ID when that ID doesn't actually exist.""" + def on_put(url, body, **kwargs): + match_data = copy.deepcopy(REPORT_INIT) + match_data['timestamp'] = body['timestamp'] + assert body == match_data + return body + + cbcsdk_mock.mock_request("PUT", "/threathunter/watchlistmgr/v3/orgs/test/reports/" + "69e2a8d0-bc36-4970-9834-8687efe1aff7", on_put) + api = cbcsdk_mock.api + report = Report(api, None, copy.deepcopy(REPORT_INIT), None, True) + report.remove_iocs_by_id(['notexist']) + report.update() + assert len(report.iocs_) == 2 + + +def test_feed_builder_save(cbcsdk_mock): + """Tests the operation of the FeedBuilder and the save() function.""" + def on_post(url, body, **kwargs): + body_data = copy.deepcopy(body) + for r in body_data['reports']: + if 'id' in r: + assert re.fullmatch(GUID_PATTERN, r['id']) + del r['id'] + assert body_data == FEED_BUILT_VIA_BUILDER + return_value = copy.deepcopy(body['feedinfo']) + return_value["id"] = "qwertyuiop" + return_value["owner"] = "JRN" + return_value["access"] = "private" + return return_value + + cbcsdk_mock.mock_request("POST", "/threathunter/feedmgr/v2/orgs/test/feeds", on_post) + api = cbcsdk_mock.api + # Start by building a report for the new feed to contain. + report_builder = Report.create(api, "NotReal", "Not real description", 2) + report_builder.set_title("ReportTitle").set_description("The report description").set_timestamp(1234567890) + report_builder.set_severity(5).set_link('https://example.com').add_tag("Alpha").add_tag("Bravo") + report_builder.add_ioc(IOC_V2.create_equality(api, "foo", "process_name", "evil.exe")) + report_builder.add_ioc(IOC_V2.create_equality(api, "bar", "netconn_ipv4", "10.29.99.1")) + report_builder.set_visibility("visible") + report = report_builder.build() + # Now build the feed. + builder = Feed.create(api, "NotReal", "http://127.0.0.1", "Not a real summary", "Fake") + builder.set_name("FeedName").set_provider_url("http://example.com").set_summary("Summary information") + builder.set_category("Intrusion").set_source_label("SourceLabel").add_reports([report]) + feed = builder.build() + feed.save() + assert feed.id == "qwertyuiop" + assert len(feed._reports) == 1 + + +@pytest.mark.parametrize("feed_init, do_operation, expectation", [ + (FEED_INIT, True, does_not_raise()), + (FEED_BUILT_VIA_BUILDER, False, pytest.raises(InvalidObjectError)) +]) +def test_delete_feed(cbcsdk_mock, feed_init, do_operation, expectation): + """Tests the delete() function.""" + if do_operation: + cbcsdk_mock.mock_request("DELETE", "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop", + CBCSDKMock.StubResponse(None, 204)) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=feed_init) + with expectation: + feed.delete() + + +@pytest.mark.parametrize("feed_init, change_item, new_value, do_operation, new_info, expectation", [ + (FEED_INIT, 'name', 'NewName', True, FEED_UPDATE_INFO_1, does_not_raise()), + (FEED_BUILT_VIA_BUILDER, 'name', 'NewName', False, None, pytest.raises(InvalidObjectError)) +]) +def test_update_feed_info(cbcsdk_mock, feed_init, change_item, new_value, do_operation, new_info, expectation): + """Tests the update() function.""" + def on_put(url, body, **kwargs): + assert body == new_info + return new_info + + if do_operation: + cbcsdk_mock.mock_request("PUT", "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop/feedinfo", on_put) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=copy.deepcopy(feed_init)) + params = {change_item: new_value} + with expectation: + feed.update(**params) + assert feed._info[change_item] == new_value + + +@pytest.mark.parametrize("feed_init, do_request, expectation", [ + (FEED_INIT, True, does_not_raise()), + (FEED_BUILT_VIA_BUILDER, False, pytest.raises(InvalidObjectError)) +]) +def test_get_feed_reports(cbcsdk_mock, feed_init, do_request, expectation): + """Tests the reports() property.""" + if do_request: + cbcsdk_mock.mock_request("GET", "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop/reports", REPORT_GET_RESP) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=feed_init) + with expectation: + reports = feed.reports + assert len(reports) == 1 + assert reports[0].id == '109027d1-064c-477d-aa34-528606ef72a9' + + +@pytest.mark.parametrize("feed_init, do_request, expectation", [ + (FEED_INIT, True, does_not_raise()), + (FEED_BUILT_VIA_BUILDER, False, pytest.raises(InvalidObjectError)) +]) +def test_replace_reports(cbcsdk_mock, feed_init, do_request, expectation): + """Tests the replace_reports method.""" + def on_post(url, body, **kwargs): + array = body['reports'] + assert len(array) == 1 + assert array[0]['id'] == "065fb68d-42a8-4b2e-8f91-17f925f54356" + return {"success": True} + + if do_request: + cbcsdk_mock.mock_request("POST", "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop/reports", on_post) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=feed_init) + new_report = Report(api, None, REPORT_INIT_2) + with expectation: + feed.replace_reports([new_report]) + + +@pytest.mark.parametrize("feed_init, do_request, expectation", [ + (FEED_INIT, True, does_not_raise()), + (FEED_BUILT_VIA_BUILDER, False, pytest.raises(InvalidObjectError)) +]) +def test_append_reports(cbcsdk_mock, feed_init, do_request, expectation): + """Tests the append_reports method.""" + def on_post(url, body, **kwargs): + array = body['reports'] + assert len(array) == 2 + assert array[0]['id'] == "69e2a8d0-bc36-4970-9834-8687efe1aff7" + assert array[1]['id'] == "065fb68d-42a8-4b2e-8f91-17f925f54356" + return {"success": True} + + if do_request: + cbcsdk_mock.mock_request("POST", "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop/reports", on_post) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=feed_init) + new_report = Report(api, None, REPORT_INIT_2) + with expectation: + feed.append_reports([new_report]) + + +def test_watchlist_builder(cbcsdk_mock): + """Tests the function of the watchlist builder.""" + def on_post(url, body, **kwargs): + assert body == WATCHLIST_BUILDER_IN + return WATCHLIST_BUILDER_OUT + + cbcsdk_mock.mock_request('POST', "/threathunter/watchlistmgr/v3/orgs/test/watchlists", on_post) + api = cbcsdk_mock.api + report = Report(api, None, REPORT_INIT) + report._from_watchlist = True + builder = Watchlist.create(api, "NameErased") + builder.set_name('NewWatchlist').set_description('I am a watchlist').set_tags_enabled(False) + builder.set_alerts_enabled(True).add_report_ids(['47474d40-1f94-4995-b6d9-1d1eea3528b3']).add_reports([report]) + watchlist = builder.build() + watchlist.save() + assert watchlist.id == "ABCDEFGHabcdefgh" + + +def test_create_watchlist_from_feed(cbcsdk_mock): + """Tests the function of create_from_feed().""" + def on_post(url, body, **kwargs): + assert body == WATCHLIST_FROM_FEED_IN + return WATCHLIST_FROM_FEED_OUT + + cbcsdk_mock.mock_request('POST', "/threathunter/watchlistmgr/v3/orgs/test/watchlists", on_post) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=FEED_INIT) + watchlist = Watchlist.create_from_feed(feed) + watchlist.save() + assert watchlist.id == "ABCDEFGHabcdefgh" + + +def test_watchlist_add_report_ids(cbcsdk_mock): + """Tests the add_report_ids() method.""" + def on_put(url, body, **kwargs): + assert body['report_ids'] == ADD_REPORT_IDS_LIST + return_value = copy.deepcopy(WATCHLIST_BUILDER_OUT) + return_value['report_ids'] = ADD_REPORT_IDS_LIST + return return_value + + cbcsdk_mock.mock_request("PUT", "/threathunter/watchlistmgr/v3/orgs/test/watchlists/ABCDEFGHabcdefgh", on_put) + api = cbcsdk_mock.api + watchlist = Watchlist(api, initial_data=copy.deepcopy(WATCHLIST_BUILDER_OUT)) + watchlist.add_report_ids(['64414d1f-66d5-4b32-9b8a-b778e13ab836', 'c9826407-0d5a-467f-bc98-7da2035de1bc']) + assert watchlist.report_ids == ADD_REPORT_IDS_LIST + + +def test_watchlist_add_reports(cbcsdk_mock): + """Tests the add_reports() method.""" + def on_put(url, body, **kwargs): + assert body['report_ids'] == ADD_REPORTS_LIST + return_value = copy.deepcopy(WATCHLIST_BUILDER_OUT) + return_value['report_ids'] = ADD_REPORTS_LIST + return return_value + + cbcsdk_mock.mock_request("PUT", "/threathunter/watchlistmgr/v3/orgs/test/watchlists/ABCDEFGHabcdefgh", on_put) + api = cbcsdk_mock.api + watchlist = Watchlist(api, initial_data=copy.deepcopy(WATCHLIST_BUILDER_OUT)) + report = Report(api, None, REPORT_INIT_2) + report._from_watchlist = True + watchlist.add_reports([report]) + assert watchlist.report_ids == ADD_REPORTS_LIST + + +@pytest.mark.parametrize("data, expectation, message", [ + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, does_not_raise(), None), + ({"title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'id'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'title'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'description'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'timestamp'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'severity'"), + ({"id": 123, "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'id' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": 6781, "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'title' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": 8, + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'description' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": "1234567890", "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'timestamp' is not an integer"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": "5", "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'severity' is not an integer"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": 8888, "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'link' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": "Alpha", + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'tags' is not a list"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": {"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}, + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report field 'iocs_v2' is not a list"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": 6681}, pytest.raises(InvalidObjectError), "Report field 'visibility' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": -123456, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Timestamp cannot be negative"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 18, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Severity value out of range"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", 16], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report tag is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'iocs_v2'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Report should have at least one IOC"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'match_type': 'query', 'values': ['foo']}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'id'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 'foo', 'values': ['foo']}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'match_type'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 'foo', 'match_type': 'query'}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "Missing key: 'values'"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 14, 'match_type': 'query', 'values': ['foo']}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "IOC field 'id' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 'foo', 'match_type': 303, 'values': ['foo']}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "IOC field 'match_type' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 'foo', 'match_type': 'query', 'values': 'blort'}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "IOC field 'values' is not a list"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 'foo', 'match_type': 'query', 'values': ['foo'], 'field': 15}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "IOC field 'field' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 'foo', 'match_type': 'query', 'values': ['foo'], 'link': 15}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "IOC field 'link' is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{'id': 'foo', 'match_type': 'notdef', 'values': ['foo']}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), + "error in IOC 'match_type' value: Invalid match type"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "values": ["evil.exe"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "IOC of type equality must have a 'field' value"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "equality", "field": "process_name", "values": ["evil.exe", 9]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "IOC value is not a string"), + ({"id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", "title": "ReportTitle", "description": "The report description", + "timestamp": 1234567890, "severity": 5, "link": "https://example.com", "tags": ["Alpha", "Bravo"], + "iocs_v2": [{"id": "foo", "match_type": "query", "values": ["process_name:evil.exe", "netconn_ipv4:10.0.0.1"]}], + "visibility": "visible"}, pytest.raises(InvalidObjectError), "query IOC should have one and only one value"), +]) +def test_feed_report_validation(data, expectation, message): + """Tests the _validate_report_rawdata method.""" + look_at_message = True + with expectation as e: + Feed._validate_report_rawdata([data]) + look_at_message = False + if look_at_message: + assert e.type is InvalidObjectError + assert e.value.args[0] == message + + +@pytest.mark.parametrize("feed_init, do_request, expectation", [ + (FEED_INIT, True, does_not_raise()), + (FEED_BUILT_VIA_BUILDER, False, pytest.raises(InvalidObjectError)) +]) +def test_replace_reports_rawdata(cbcsdk_mock, feed_init, do_request, expectation): + """Tests the replace_reports_rawdata method.""" + def on_post(url, body, **kwargs): + array = body['reports'] + assert len(array) == 1 + assert array[0]['id'] == "065fb68d-42a8-4b2e-8f91-17f925f54356" + return {"success": True} + + if do_request: + cbcsdk_mock.mock_request("POST", "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop/reports", on_post) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=feed_init) + with expectation: + feed.replace_reports_rawdata([REPORT_INIT_2]) + + +@pytest.mark.parametrize("feed_init, do_request, expectation", [ + (FEED_INIT, True, does_not_raise()), + (FEED_BUILT_VIA_BUILDER, False, pytest.raises(InvalidObjectError)) +]) +def test_append_reports_rawdata(cbcsdk_mock, feed_init, do_request, expectation): + """Tests the append_reports_rawdata method.""" + def on_post(url, body, **kwargs): + array = body['reports'] + assert len(array) == 2 + assert array[0]['id'] == "69e2a8d0-bc36-4970-9834-8687efe1aff7" + assert array[1]['id'] == "065fb68d-42a8-4b2e-8f91-17f925f54356" + return {"success": True} + + if do_request: + cbcsdk_mock.mock_request("POST", "/threathunter/feedmgr/v2/orgs/test/feeds/qwertyuiop/reports", on_post) + api = cbcsdk_mock.api + feed = Feed(api, initial_data=feed_init) + with expectation: + feed.append_reports_rawdata([REPORT_INIT_2]) diff --git a/src/tests/unit/enterprise_edr/test_enterprise_edr_ubs.py b/src/tests/unit/enterprise_edr/test_enterprise_edr_ubs.py index 79e4189e3..43cbcc1ca 100644 --- a/src/tests/unit/enterprise_edr/test_enterprise_edr_ubs.py +++ b/src/tests/unit/enterprise_edr/test_enterprise_edr_ubs.py @@ -62,6 +62,34 @@ def post_validate(url, body, **kwargs): assert url is not None +def test_binary_query_case_insensitive(cbcsdk_mock): + """Testing Binary Querying""" + called = False + + def post_validate(url, body, **kwargs): + nonlocal called + if not called: + called = True + assert body['expiration_seconds'] == 3600 + else: + assert body['expiration_seconds'] == 10 + return BINARY_GET_FILE_RESP + + sha256 = "00A16C806FF694B64E566886BBA5122655EFF89B45226CDDC8651DF7860E4524" + cbcsdk_mock.mock_request("GET", f"/ubs/v1/orgs/test/sha256/{sha256}", BINARY_GET_METADATA_RESP) + api = cbcsdk_mock.api + binary = api.select(Binary, sha256) + assert isinstance(binary, Binary) + cbcsdk_mock.mock_request("GET", f"/ubs/v1/orgs/test/sha256/{sha256}/summary/device", BINARY_GET_DEVICE_SUMMARY_RESP) + summary = binary.summary + cbcsdk_mock.mock_request("POST", "/ubs/v1/orgs/test/file/_download", post_validate) + url = binary.download_url() + assert summary is not None + assert url is not None + url = binary.download_url(expiration_seconds=10) + assert url is not None + + def test_binary_query_error(cbcsdk_mock): """Testing Binary Querying Error""" api = cbcsdk_mock.api diff --git a/src/tests/unit/fixtures/enterprise_edr/mock_threatintel.py b/src/tests/unit/fixtures/enterprise_edr/mock_threatintel.py index 42ba07735..a440c6f2a 100644 --- a/src/tests/unit/fixtures/enterprise_edr/mock_threatintel.py +++ b/src/tests/unit/fixtures/enterprise_edr/mock_threatintel.py @@ -1002,3 +1002,240 @@ } ] } + +IOC_GET_IGNORED = { + "ignored": True +} + +REPORT_BUILT_VIA_BUILDER = { + "title": "ReportTitle", + "description": "The report description", + "timestamp": 1234567890, + "severity": 5, + "link": "https://example.com", + "tags": ["Alpha", "Bravo"], + "iocs_v2": [ + { + "id": "foo", + "match_type": "equality", + "field": "process_name", + "values": ["evil.exe"], + }, + { + "id": "bar", + "match_type": "equality", + "field": "netconn_ipv4", + "values": ["10.29.99.1"], + } + ], + "visibility": "visible" +} + +REPORT_INIT = { + "id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", + "title": "ReportTitle", + "description": "The report description", + "timestamp": 1234567890, + "severity": 5, + "link": "https://example.com", + "tags": ["Alpha", "Bravo"], + "iocs_v2": [ + { + "id": "foo", + "match_type": "equality", + "field": "process_name", + "values": ["evil.exe"], + }, + { + "id": "bar", + "match_type": "equality", + "field": "netconn_ipv4", + "values": ["10.29.99.1"], + } + ], + "visibility": "visible" +} + +REPORT_GET_IGNORED = { + "ignored": True +} + +REPORT_GET_SEVERITY = { + "report_id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", + "severity": 8 +} + +REPORT_UPDATE_AFTER_ADD_IOC = { + "id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", + "title": "ReportTitle", + "description": "The report description", + "timestamp": 1234567890, + "severity": 5, + "link": "https://example.com", + "tags": ["Alpha", "Bravo"], + "iocs_v2": [ + { + "id": "foo", + "match_type": "equality", + "field": "process_name", + "values": ["evil.exe"] + }, + { + "id": "bar", + "match_type": "equality", + "field": "netconn_ipv4", + "values": ["10.29.99.1"] + }, + { + "id": "quux", + "match_type": "query", + "values": ["filemod_name: \"audio.dat\""] + } + ], + "visibility": "visible" +} + +REPORT_UPDATE_AFTER_REMOVE_IOC = { + "id": "69e2a8d0-bc36-4970-9834-8687efe1aff7", + "title": "ReportTitle", + "description": "The report description", + "timestamp": 1234567890, + "severity": 5, + "link": "https://example.com", + "tags": ["Alpha", "Bravo"], + "iocs_v2": [ + { + "id": "bar", + "match_type": "equality", + "field": "netconn_ipv4", + "values": ["10.29.99.1"], + } + ], + "visibility": "visible" +} + +FEED_BUILT_VIA_BUILDER = { + "feedinfo": { + "name": "FeedName", + "provider_url": "http://example.com", + "summary": "Summary information", + "category": "Intrusion", + "source_label": "SourceLabel" + }, + "reports": [REPORT_BUILT_VIA_BUILDER] +} + +FEED_INIT = { + "feedinfo": { + "id": "qwertyuiop", + "name": "FeedName", + "provider_url": "http://example.com", + "summary": "Summary information", + "category": "Intrusion", + "source_label": "SourceLabel", + "owner": "JRN", + "access": "private" + }, + "reports": [REPORT_INIT] +} + +FEED_UPDATE_INFO_1 = { + "id": "qwertyuiop", + "name": "NewName", + "provider_url": "http://example.com", + "summary": "Summary information", + "category": "Intrusion", + "source_label": "SourceLabel", + "owner": "JRN", + "access": "private" +} + +REPORT_INIT_2 = { + "id": "065fb68d-42a8-4b2e-8f91-17f925f54356", + "title": "SecondTitle", + "description": "The second description", + "timestamp": 1234567890, + "severity": 4, + "link": "https://example.com", + "tags": ["Alpha"], + "iocs_v2": [ + { + "id": "1foo", + "match_type": "equality", + "field": "process_name", + "values": ["wicked.exe"], + }, + { + "id": "1bar", + "match_type": "equality", + "field": "netconn_ipv4", + "values": ["10.29.99.64"], + } + ], + "visibility": "visible" +} + +WATCHLIST_BUILDER_IN = { + "name": "NewWatchlist", + "description": "I am a watchlist", + "tags_enabled": False, + "alerts_enabled": True, + "report_ids": [ + '47474d40-1f94-4995-b6d9-1d1eea3528b3', + "69e2a8d0-bc36-4970-9834-8687efe1aff7" + ] +} + +WATCHLIST_BUILDER_OUT = { + "name": "NewWatchlist", + "description": "I am a watchlist", + "id": "ABCDEFGHabcdefgh", + "tags_enabled": False, + "alerts_enabled": True, + "create_timestamp": 1234567890, + "last_update_timestamp": 1234567890, + "classifier": None, + "report_ids": [ + '47474d40-1f94-4995-b6d9-1d1eea3528b3', + "69e2a8d0-bc36-4970-9834-8687efe1aff7" + ] +} + +WATCHLIST_FROM_FEED_IN = { + "name": "Feed FeedName", + "description": "Summary information", + "tags_enabled": True, + "alerts_enabled": False, + "classifier": { + "key": "feed_id", + "value": "qwertyuiop" + } +} + +WATCHLIST_FROM_FEED_OUT = { + "name": "Feed FeedName", + "description": "Summary information", + "id": "ABCDEFGHabcdefgh", + "tags_enabled": True, + "alerts_enabled": False, + "create_timestamp": 1234567890, + "last_update_timestamp": 1234567890, + "report_ids": None, + "classifier": { + "key": "feed_id", + "value": "qwertyuiop" + } +} + +ADD_REPORT_IDS_LIST = [ + '47474d40-1f94-4995-b6d9-1d1eea3528b3', + "69e2a8d0-bc36-4970-9834-8687efe1aff7", + '64414d1f-66d5-4b32-9b8a-b778e13ab836', + 'c9826407-0d5a-467f-bc98-7da2035de1bc' +] + +ADD_REPORTS_LIST = [ + '47474d40-1f94-4995-b6d9-1d1eea3528b3', + "69e2a8d0-bc36-4970-9834-8687efe1aff7", + "065fb68d-42a8-4b2e-8f91-17f925f54356" +] diff --git a/src/tests/unit/fixtures/live_response/mock_command.py b/src/tests/unit/fixtures/live_response/mock_command.py index c232482a3..3668118f0 100755 --- a/src/tests/unit/fixtures/live_response/mock_command.py +++ b/src/tests/unit/fixtures/live_response/mock_command.py @@ -96,6 +96,17 @@ 'completion_time': 2345678901 } +GET_FILE_CANCELLED_RESP = { + 'id': 7, + 'session_id': '1:2468', + 'device_id': 2468, + 'command_timeout': 120, + 'status': 'CANCELLED', + 'name': 'get file', + 'object': 'C:\\\\test.txt', + 'completion_time': 2345678901 +} + GET_FILE_COMMAND_RESP = { "values": [], "file_details": { diff --git a/src/tests/unit/fixtures/platform/mock_process.py b/src/tests/unit/fixtures/platform/mock_process.py index 113ff837f..9adb5628a 100644 --- a/src/tests/unit/fixtures/platform/mock_process.py +++ b/src/tests/unit/fixtures/platform/mock_process.py @@ -2446,6 +2446,14 @@ ] } +GET_PROCESS_DETAILS_JOB_RESULTS_RESP_ZERO = { + 'contacted': 0, + 'completed': 0, + 'num_available': 0, + 'num_found': 0, + 'results': [] +} + GET_FACET_SEARCH_RESULTS_RESP = { "ranges": [ { diff --git a/src/tests/unit/platform/test_platform_models.py b/src/tests/unit/platform/test_platform_models.py index 789a06f25..89e41d27e 100755 --- a/src/tests/unit/platform/test_platform_models.py +++ b/src/tests/unit/platform/test_platform_models.py @@ -29,7 +29,7 @@ def __init__(self, expected_id): self.expected_id = expected_id self.was_called = False - def request_session(self, sensor_id): + def request_session(self, sensor_id, async_mode=False): """ Stub out the request_session call to the scheduler. diff --git a/src/tests/unit/platform/test_platform_process.py b/src/tests/unit/platform/test_platform_process.py index 390a02e3b..1344f017a 100644 --- a/src/tests/unit/platform/test_platform_process.py +++ b/src/tests/unit/platform/test_platform_process.py @@ -40,7 +40,8 @@ EXPECTED_PROCESS_FACETS, EXPECTED_PROCESS_RANGES_FACETS, GET_PROCESS_TREE_STR, - GET_PROCESS_SUMMARY_STR) + GET_PROCESS_SUMMARY_STR, + GET_PROCESS_DETAILS_JOB_RESULTS_RESP_ZERO) log = logging.basicConfig(format='%(asctime)s %(levelname)s:%(message)s', level=logging.DEBUG, filename='log.txt') @@ -1106,6 +1107,25 @@ def test_process_get_details(cbcsdk_mock): assert 10222 in results['process_pid'] +def test_process_get_details_zero(cbcsdk_mock): + """Test get_details on a process.""" + cbcsdk_mock.mock_request("POST", "/api/investigate/v2/orgs/test/processes/detail_jobs", + POST_PROCESS_DETAILS_JOB_RESP) + cbcsdk_mock.mock_request("GET", + "/api/investigate/v2/orgs/test/processes/detail_jobs/ccc47a52-9a61-4c77-8652-8a03dc187b98", # noqa: E501 + GET_PROCESS_DETAILS_JOB_STATUS_RESP) + cbcsdk_mock.mock_request("GET", + "/api/investigate/v2/orgs/test/processes/detail_jobs/ccc47a52-9a61-4c77-8652-8a03dc187b98/results", # noqa: E501 + GET_PROCESS_DETAILS_JOB_RESULTS_RESP_ZERO) + + api = cbcsdk_mock.api + process = Process(api, '80dab519-3b5f-4502-afad-da87cd58a4c3', + {'process_guid': '80dab519-3b5f-4502-afad-da87cd58a4c3'}) + results = process.get_details() + assert results['process_guid'] == '80dab519-3b5f-4502-afad-da87cd58a4c3' + assert results.get('device_id') is None + + def test_process_get_details_async(cbcsdk_mock): """Test get_details on a process in async mode.""" cbcsdk_mock.mock_request("POST", "/api/investigate/v2/orgs/test/processes/detail_jobs", diff --git a/src/tests/unit/test_live_response_api.py b/src/tests/unit/test_live_response_api.py index d85307efe..0362074f9 100755 --- a/src/tests/unit/test_live_response_api.py +++ b/src/tests/unit/test_live_response_api.py @@ -63,7 +63,8 @@ MEMDUMP_START_RESP, MEMDUMP_END_RESP, MEMDUMP_DEL_START_RESP, - MEMDUMP_DEL_END_RESP) + MEMDUMP_DEL_END_RESP, + GET_FILE_CANCELLED_RESP) from tests.unit.fixtures.live_response.mock_device import DEVICE_RESPONSE, UDEVICE_RESPONSE, POST_DEVICE_SEARCH_RESP from tests.unit.fixtures.live_response.mock_session import (SESSION_INIT_RESP, SESSION_POLL_RESP, SESSION_POLL_RESP_ERROR, USESSION_INIT_RESP, @@ -245,6 +246,23 @@ def test_create_session(cbcsdk_mock): assert session.os_type == 1 +def test_create_session_async(cbcsdk_mock): + """Test creating a Live Response session.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + session_id, result = manager.request_session(2468, async_mode=True) + session = result.result() + assert session.session_id == '1:2468' + assert session.device_id == 2468 + assert session._cblr_manager is manager + assert session._cb is cbcsdk_mock.api + assert session.os_type == 1 + session.close() + + def test_create_session_with_poll_error(cbcsdk_mock): """Test creating a Live Response session with an error in the polling.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -294,6 +312,34 @@ def test_create_session_with_keepalive_option(cbcsdk_mock): manager.stop_keepalive_thread() +def test_create_session_with_keepalive_option_async(cbcsdk_mock): + """Test creating a Live Response session using the keepalive option.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api, 100000, True) + try: + session_id, result = manager.request_session(2468, async_mode=True) + session1 = result.result() + assert session_id == '1:2468' + assert session1.session_id == '1:2468' + assert session1.device_id == 2468 + assert session1._cblr_manager is manager + assert session1._cb is cbcsdk_mock.api + assert session1.os_type == 1 + session1.close() + session_id, result2 = manager.request_session(2468, async_mode=True) + session2 = result2.result() + assert session2 is session1 + session2.close() + assert len(manager._sessions) == 1 + manager._maintain_sessions() + assert len(manager._sessions) == 0 + finally: + manager.stop_keepalive_thread() + + @pytest.mark.parametrize("thrown_exception", [ (ObjectNotFoundError('/appservices/v6/orgs/test/liveresponse/sessions/1:2468'),), (ServerError(404, 'test error'),) @@ -337,6 +383,29 @@ def test_list_directory(cbcsdk_mock): assert 'ARCHIVE' in files[2]['attributes'] +def test_list_directory_async(cbcsdk_mock): + """Test the response to the 'list directory' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + DIRECTORY_LIST_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/6', + DIRECTORY_LIST_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, result = session.list_directory('C:\\\\TEMP\\\\', async_mode=True) + assert command_id == 6 + files = result.result() + assert files[0]['filename'] == '.' + assert 'DIRECTORY' in files[0]['attributes'] + assert files[1]['filename'] == '..' + assert 'DIRECTORY' in files[1]['attributes'] + assert files[2]['filename'] == 'test.txt' + assert 'ARCHIVE' in files[2]['attributes'] + + def test_delete_file(cbcsdk_mock): """Test the response to the 'delete file' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -352,6 +421,22 @@ def test_delete_file(cbcsdk_mock): session.delete_file('C:\\\\TEMP\\\\foo.txt') +def test_delete_file_async(cbcsdk_mock): + """Test the response to the 'delete file' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + DELETE_FILE_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/3', + DELETE_FILE_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, _ = session.delete_file('C:\\\\TEMP\\\\foo.txt', async_mode=True) + assert command_id == 3 + + def test_delete_file_with_error(cbcsdk_mock): """Test the response to the 'delete file' command when it returns an error.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -370,7 +455,7 @@ def test_delete_file_with_error(cbcsdk_mock): def test_get_file(cbcsdk_mock, connection_mock): - """Test the response to the 'delete file' command.""" + """Test the response to the 'get file' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) @@ -384,6 +469,137 @@ def test_get_file(cbcsdk_mock, connection_mock): session.get_file('c:\\\\test.txt') +def test_get_file_cancelled(cbcsdk_mock, connection_mock): + """Test the response to the 'get file' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + GET_FILE_COMMAND_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + GET_FILE_CANCELLED_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + with pytest.raises(ApiError) as ex: + session.get_file('c:\\\\test.txt') + assert 'The command has been cancelled.' in str(ex.value) + + +def test_get_file_cancelled_async(cbcsdk_mock, connection_mock): + """Test the response to the 'get file' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + GET_FILE_COMMAND_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + GET_FILE_CANCELLED_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + with pytest.raises(ApiError) as ex: + _, result = session.get_file('c:\\\\test.txt', async_mode=True) + result.result() + assert 'The command has been cancelled.' in str(ex.value) + + +def test_get_file_async(cbcsdk_mock, connection_mock): + """Test the response to the 'get file' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + GET_FILE_COMMAND_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + GET_FILE_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, _ = session.get_file('c:\\\\test.txt', async_mode=True) + assert command_id == 7 + + +def test_get_raw_file_async(cbcsdk_mock, connection_mock): + """Test the response to the 'get file' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + GET_FILE_COMMAND_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + GET_FILE_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, _ = session.get_raw_file('c:\\\\test.txt', async_mode=True) + assert command_id == 7 + + +def test_command_status(cbcsdk_mock): + """Test command status method""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + GET_FILE_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + status = session.command_status(7) + assert status == 'COMPLETE' + + +def test_session_status(cbcsdk_mock): + """Test command status method""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + session_id, _ = manager.request_session(2468, async_mode=True) + status = manager.session_status(session_id) + assert status == 'ACTIVE' + + +def test_cancel_complete_command(cbcsdk_mock): + """Test the response to the 'cancel command' command for completed command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + GET_FILE_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + with pytest.raises(ApiError) as excinfo: + session.cancel_command(7) + assert excinfo.value.__str__().startswith('Cannot cancel command in status COMPLETE') + + +def test_cancel_pending_command(cbcsdk_mock): + """Test the response to the 'cancel command' command for completed command.""" + _was_called = False + + def delete_req(url, body): + nonlocal _was_called + _was_called = True + return None + + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + GET_FILE_COMMAND_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', delete_req) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + try: + session.cancel_command(7) + assert _was_called + except ApiError: + raise Exception('Failed') + + def test_put_file(cbcsdk_mock, mox): """Test the response to the 'put file' command.""" def respond_to_post(url, body, **kwargs): @@ -410,6 +626,33 @@ def respond_to_post(url, body, **kwargs): mox.VerifyAll() +def test_put_file_async(cbcsdk_mock, mox): + """Test the response to the 'put file' command.""" + def respond_to_post(url, body, **kwargs): + assert body['session_id'] == '1:2468' + assert body['name'] == 'put file' + assert body['file_id'] == 10203 + assert body['path'] == 'foobar.txt' + return PUT_FILE_START_RESP + + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', respond_to_post) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/6', + PUT_FILE_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + filep = io.StringIO('This is a test') + with manager.request_session(2468) as session: + mox.StubOutWithMock(session, '_upload_file') + session._upload_file(filep).AndReturn(10203) + mox.ReplayAll() + command_id, _ = session.put_file(filep, 'foobar.txt', async_mode=True) + assert command_id == 6 + mox.VerifyAll() + + def test_create_directory(cbcsdk_mock): """Test the response to the 'create directory' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -425,6 +668,22 @@ def test_create_directory(cbcsdk_mock): session.create_directory('C:\\\\TEMP\\\\TRASH') +def test_create_directory_async(cbcsdk_mock): + """Test the response to the 'create directory' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + CREATE_DIRECTORY_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/7', + CREATE_DIRECTORY_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, _ = session.create_directory('C:\\\\TEMP\\\\TRASH', async_mode=True) + assert command_id == 7 + + def test_walk(cbcsdk_mock, mox): """Test the logic of the directory walking.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -520,6 +779,23 @@ def test_kill_process(cbcsdk_mock): assert session.kill_process(601) +def test_kill_process_async(cbcsdk_mock): + """Test the response to the 'kill' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + KILL_PROC_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/13', + KILL_PROC_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, result = session.kill_process(601, async_mode=True) + assert command_id == 13 + assert result.result() + + def test_kill_process_timeout(cbcsdk_mock): """Test the response to the 'kill' command when it times out.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -550,6 +826,23 @@ def test_create_process(cbcsdk_mock): assert session.create_process('start_daemon', False) is None +def test_create_process_async(cbcsdk_mock): + """Test the response to the 'create process' command with wait for completion.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + CREATE_PROC_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/52', + CREATE_PROC_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, result = session.create_process('start_daemon', False, async_mode=True) + assert command_id == 52 + assert result.result() is None + + def test_spawn_process(cbcsdk_mock): """Test the response to the 'create process' command without wait for completion.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -655,6 +948,27 @@ def test_list_processes(cbcsdk_mock): assert plist[2]['process_path'] == 'borg' +def test_list_processes_async(cbcsdk_mock): + """Test the response to the 'list processes' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + LIST_PROC_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/10', + LIST_PROC_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + command_id, result = session.list_processes(async_mode=True) + plist = result.result() + assert command_id == 10 + assert len(plist) == 3 + assert plist[0]['process_path'] == 'proc1' + assert plist[1]['process_path'] == 'server' + assert plist[2]['process_path'] == 'borg' + + def test_registry_enum(cbcsdk_mock): """Test the response to the 'reg enum keys' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -681,6 +995,38 @@ def test_registry_enum(cbcsdk_mock): assert keyitem['value_name'] in value_names +def test_registry_enum_async(cbcsdk_mock): + """Test the response to the 'reg enum keys' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + REG_ENUM_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/56', + REG_ENUM_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + c_id, result = session.list_registry_keys_and_values('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI', + async_mode=True) + rc1 = result.result() + assert c_id == 56 + assert len(rc1['sub_keys']) == 2 + assert 'Parameters' in rc1['sub_keys'] + assert 'Enum' in rc1['sub_keys'] + value_names = ['Start', 'Type', 'ErrorControl', 'ImagePath', 'DisplayName', 'Group', 'DriverPackageId', 'Tag'] + assert len(rc1['values']) == len(value_names) + for keyitem in rc1['values']: + assert keyitem['value_name'] in value_names + command_id, result = session.list_registry_values('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI', + async_mode=True) + assert command_id == 56 + rc2 = result.result() + assert len(rc2) == len(value_names) + for keyitem in rc2: + assert keyitem['value_name'] in value_names + + def test_registry_get(cbcsdk_mock): """Test the response to the 'reg get value' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -699,6 +1045,27 @@ def test_registry_get(cbcsdk_mock): assert val['value_type'] == 'REG_DWORD' +def test_registry_get_async(cbcsdk_mock): + """Test the response to the 'reg get value' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + REG_GET_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/61', + REG_GET_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + c_id, result = session.get_registry_value('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\Start', + async_mode=True) + val = result.result() + assert c_id == 61 + assert val['value_data'] == 0 + assert val['value_name'] == 'Start' + assert val['value_type'] == 'REG_DWORD' + + @pytest.mark.parametrize("set_val,check_val,overwrite,set_type,check_type", [ (42, 42, False, None, 'REG_DWORD'), (['a', 'b', 'c'], ['a', 'b', 'c'], True, None, 'REG_MULTI_SZ'), @@ -730,6 +1097,41 @@ def respond_to_post(url, body, **kwargs): overwrite, set_type) +@pytest.mark.parametrize("set_val,check_val,overwrite,set_type,check_type", [ + (42, 42, False, None, 'REG_DWORD'), + (['a', 'b', 'c'], ['a', 'b', 'c'], True, None, 'REG_MULTI_SZ'), + ([10, 20, 30], ['10', '20', '30'], False, None, 'REG_MULTI_SZ'), + ('Quimby', 'Quimby', True, None, 'REG_SZ'), + (80231, 80231, False, 'REG_QWORD', 'REG_QWORD') +]) +def test_registry_set_async(cbcsdk_mock, set_val, check_val, overwrite, set_type, check_type): + """Test the response to the 'reg set value' command.""" + def respond_to_post(url, body, **kwargs): + assert body['session_id'] == '1:2468' + assert body['name'] == 'reg set value' + assert body['path'] == 'HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\testvalue' + assert body['overwrite'] == overwrite + assert body['value_type'] == check_type + assert body['value_data'] == check_val + return REG_SET_START_RESP + + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', respond_to_post) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/62', + REG_SET_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + c_id, _ = session.set_registry_value('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\testvalue', + set_val, + overwrite, + set_type, + async_mode=True) + assert c_id == 62 + + def test_registry_create_key(cbcsdk_mock): """Test the response to the 'reg create key' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -745,6 +1147,23 @@ def test_registry_create_key(cbcsdk_mock): session.create_registry_key('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\Nonsense') +def test_registry_create_key_async(cbcsdk_mock): + """Test the response to the 'reg create key' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + REG_CREATE_KEY_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/63', + REG_CREATE_KEY_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + c_id, _ = session.create_registry_key('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\Nonsense', + async_mode=True) + assert c_id == 63 + + def test_registry_delete_key(cbcsdk_mock): """Test the response to the 'reg delete key' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -760,6 +1179,23 @@ def test_registry_delete_key(cbcsdk_mock): session.delete_registry_key('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\Nonsense') +def test_registry_delete_key_async(cbcsdk_mock): + """Test the response to the 'reg delete key' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + REG_DELETE_KEY_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/64', + REG_DELETE_KEY_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + c_id, _ = session.delete_registry_key('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\Nonsense', + async_mode=True) + assert c_id == 64 + + def test_registry_delete(cbcsdk_mock): """Test the response to the 'reg delete value' command.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) @@ -775,6 +1211,23 @@ def test_registry_delete(cbcsdk_mock): session.delete_registry_value('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\testvalue') +def test_registry_delete_async(cbcsdk_mock): + """Test the response to the 'reg delete value' command.""" + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', + REG_DELETE_START_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/65', + REG_DELETE_END_RESP) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + c_id, _ = session.delete_registry_value('HKLM\\SYSTEM\\CurrentControlSet\\services\\ACPI\\testvalue', + async_mode=True) + assert c_id == 65 + + def test_registry_unsupported_command(cbcsdk_mock): """Test the response to a command that we know isn't supported on the target node.""" cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', USESSION_INIT_RESP) @@ -840,6 +1293,56 @@ def respond_delete_file(url, query_parameters, default): memdump.delete() +def test_memdump_async(cbcsdk_mock): + """Test the response to the 'memdump' command.""" + generated_file_name = None + target_file_name = None + + def respond_to_post(url, body, **kwargs): + assert body['session_id'] == '1:2468' + nonlocal generated_file_name, target_file_name + if body['name'] == 'memdump': + generated_file_name = body['path'] + target_file_name = generated_file_name + print(target_file_name) + if body['compress']: + target_file_name += '.zip' + retval = copy.deepcopy(MEMDUMP_START_RESP) + retval['path'] = generated_file_name + return retval + elif body['name'] == 'delete file': + assert body['path'] == target_file_name + retval = copy.deepcopy(MEMDUMP_DEL_START_RESP) + retval['path'] = target_file_name + return retval + else: + pytest.fail(f"Invalid command name seen: {body['name']}") + + def respond_get_memdump_file(url, query_parameters, default): + retval = copy.deepcopy(MEMDUMP_END_RESP) + retval['path'] = generated_file_name + return retval + + def respond_delete_file(url, query_parameters, default): + retval = copy.deepcopy(MEMDUMP_DEL_END_RESP) + retval['path'] = target_file_name + return retval + + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions', SESSION_INIT_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', SESSION_POLL_RESP) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/devices/2468', DEVICE_RESPONSE) + cbcsdk_mock.mock_request('POST', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands', respond_to_post) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/101', + respond_get_memdump_file) + cbcsdk_mock.mock_request('GET', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468/commands/102', + respond_delete_file) + cbcsdk_mock.mock_request('DELETE', '/appservices/v6/orgs/test/liveresponse/sessions/1:2468', None) + manager = LiveResponseSessionManager(cbcsdk_mock.api) + with manager.request_session(2468) as session: + c_id, _ = session.memdump('test.txt', async_mode=True) + assert c_id == 101 + + def test_memdump_errors(cbcsdk_mock): """Test the response to the 'memdump' command.""" generated_file_name = None