def _init_snapshot(self, upload, name=None, overwrite=False,
background=False,
extra_args=None):
# type: (Union[str, IO], Optional[str], bool, bool, Optional[Dict[str, Any]]) -> Union[str, Dict[str, str]]
if self.network is None:
self.set_network()
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in self.list_snapshots():
if overwrite:
self.delete_snapshot(name)
else:
raise ValueError(
'A snapshot named ''{}'' already exists in network ''{}''. '
'Use overwrite = True if you want to overwrite the '
'existing snapshot'.format(name, self.network))
if isinstance(upload, six.string_types):
self.__init_snapshot_from_file(name, upload)
else:
if not zipfile.is_zipfile(upload):
raise ValueError("The provided data is not a valid zip file")
# upload is an IO-like object already
self.__init_snapshot_from_io(name, upload)
return self._parse_snapshot(name, background, extra_args)
def _fork_snapshot(
self,
base_name,
name=None,
overwrite=False,
background=False,
deactivate_interfaces=None,
deactivate_nodes=None,
restore_interfaces=None,
restore_nodes=None,
add_files=None,
extra_args=None,
):
# type: (str, Optional[str], bool, bool, Optional[List[Interface]], Optional[List[str]], Optional[List[Interface]], Optional[List[str]], Optional[str], Optional[Dict[str, Any]]) -> Union[str, Dict, None]
self._check_network()
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in self.list_snapshots():
if overwrite:
self.delete_snapshot(name)
else:
raise ValueError(
"A snapshot named "
"{}"
" already exists in network "
"{}"
". "
"Use overwrite = True if you want to overwrite the "
"existing snapshot".format(name, self.network)
)
encoded_file = None
if add_files is not None:
file_to_send = add_files
if os.path.isdir(add_files):
temp_zip_file = tempfile.NamedTemporaryFile()
zip_dir(add_files, temp_zip_file)
file_to_send = temp_zip_file.name
if os.path.isfile(file_to_send):
with open(file_to_send, "rb") as f:
encoded_file = base64.b64encode(f.read()).decode("ascii")
json_data = {
"snapshotBase": base_name,
"snapshotNew": name,
"deactivateInterfaces": deactivate_interfaces,
"deactivateNodes": deactivate_nodes,
"restoreInterfaces": restore_interfaces,
"restoreNodes": restore_nodes,
"zipFile": encoded_file,
}
restv2helper.fork_snapshot(self, json_data)
return self._parse_snapshot(name, background, extra_args)
def test_validate_name():
assert validate_name('goodname')
for name in [42, 'x' * 151, '/', '/etc', 'seTTings', 'settings', ".",
"../../../../etc/"]:
with pytest.raises(ValueError):
assert validate_name(name)
for name in ["23", '___', 'foo-bar']:
validate_name(name)
def bf_init_snapshot(upload,
name=None,
overwrite=False,
background=False,
extra_args=None):
# type: (str, Optional[str], bool, bool, Optional[Dict[str, Any]]) -> Union[str, Dict[str, str]]
"""Initialize a new snapshot.
:param upload: snapshot to upload
:type upload: zip file or directory
:param name: name of the snapshot to initialize
:type name: string
:param overwrite: whether or not to overwrite an existing snapshot with the
same name
:type overwrite: bool
:param background: whether or not to run the task in the background
:type background: bool
:param extra_args: extra arguments to be passed to the parse command. See bf_session.additionalArgs.
:type extra_args: dict
:return: name of initialized snapshot, or JSON dictionary of task status if background=True
:rtype: Union[str, Dict]
"""
if bf_session.network is None:
bf_set_network()
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in bf_list_snapshots():
if overwrite:
bf_delete_snapshot(name)
else:
raise ValueError('A snapshot named '
'{}'
' already exists in network '
'{}'
''.format(name, bf_session.network))
file_to_send = upload
if os.path.isdir(upload):
temp_zip_file = tempfile.NamedTemporaryFile()
zip_dir(upload, temp_zip_file)
file_to_send = temp_zip_file.name
json_data = workhelper.get_data_upload_snapshot(bf_session, name,
file_to_send)
resthelper.get_json_response(bf_session,
CoordConsts.SVC_RSC_UPLOAD_SNAPSHOT,
json_data)
return _parse_snapshot(name, background, extra_args)
def _init_snapshot(self,
upload,
name=None,
overwrite=False,
background=False,
extra_args=None):
# type: (str, Optional[str], bool, bool, Optional[Dict[str, Any]]) -> Union[str, Dict[str, str]]
if self.network is None:
self.set_network()
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in self.list_snapshots():
if overwrite:
self.delete_snapshot(name)
else:
raise ValueError(
'A snapshot named '
'{}'
' already exists in network '
'{}'
'. '
'Use overwrite = True if you want to overwrite the '
'existing snapshot'.format(name, self.network))
file_to_send = upload
tmp_file_name = None # type: Optional[Text]
if os.path.isdir(upload):
# delete=False because we re-open for reading
with tempfile.NamedTemporaryFile(delete=False) as temp_zip_file:
zip_dir(upload, temp_zip_file)
tmp_file_name = file_to_send = temp_zip_file.name
with open(file_to_send, 'rb') as fd:
json_data = workhelper.get_data_upload_snapshot(self, name, fd)
resthelper.get_json_response(self,
CoordConsts.SVC_RSC_UPLOAD_SNAPSHOT,
json_data)
# Cleanup tmp file if we made one
if tmp_file_name is not None:
try:
os.remove(tmp_file_name)
except (OSError, IOError):
# If we can't delete the file for some reason, let it be,
# no need to crash initialization
pass
return self._parse_snapshot(name, background, extra_args)
def test_validate_name():
assert validate_name("goodname")
for name in [
42,
"x" * 151,
"/",
"/etc",
"seTTings",
"settings",
".",
"../../../../etc/",
]:
with pytest.raises(ValueError):
assert validate_name(name)
for name in ["23", "___", "foo-bar"]:
validate_name(name)
def _bf_init_snapshot(upload, name, overwrite, background):
if bf_session.network is None:
bf_set_network()
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in bf_list_snapshots():
if overwrite:
bf_delete_snapshot(name)
else:
raise ValueError('A snapshot named '
'{}'
' already exists in network '
'{}'
''.format(name, bf_session.network))
file_to_send = upload
if os.path.isdir(upload):
tempFile = tempfile.NamedTemporaryFile()
zip_dir(upload, tempFile)
file_to_send = tempFile.name
json_data = workhelper.get_data_upload_snapshot(bf_session, name,
file_to_send)
resthelper.get_json_response(bf_session,
CoordConsts.SVC_RSC_UPLOAD_SNAPSHOT,
json_data)
bf_session.baseSnapshot = name
work_item = workhelper.get_workitem_parse(bf_session, name)
parse_execute = workhelper.execute(work_item,
bf_session,
background=background)
if not background:
status = parse_execute["status"]
if WorkStatusCode(status) != WorkStatusCode.TERMINATEDNORMALLY:
bf_session.baseSnapshot = None
bf_logger.info("Default snapshot is now unset")
else:
bf_logger.info("Default snapshot is now set to %s",
bf_session.baseSnapshot)
return parse_execute
def set_network(
self, name: Optional[str] = None, prefix: str = Options.default_network_prefix
) -> str:
"""
Configure the network used for analysis.
:param name: name of the network to set. If `None`, a name will be generated
:type name: str
:param prefix: prefix to prepend to auto-generated network names if name is empty
:type name: str
:return: name of the configured network
:rtype: str
:raises BatfishException: if configuration fails
"""
if name is None:
name = prefix + get_uuid()
validate_name(name, "network")
try:
net = restv2helper.get_network(self, name)
self.network = str(net["name"])
return self.network
except HTTPError as e:
if e.response.status_code != 404:
raise BatfishException("Unknown error accessing network", e)
json_data = workhelper.get_data_init_network(self, name)
json_response = resthelper.get_json_response(
self, CoordConsts.SVC_RSC_INIT_NETWORK, json_data
)
network_name = json_response.get(CoordConsts.SVC_KEY_NETWORK_NAME)
if network_name is None:
raise BatfishException(
"Network initialization failed. Server response: {}".format(
json_response
)
)
self.network = str(network_name)
return self.network
def bf_set_network(name=None, prefix=Options.default_network_prefix):
# type: (str, str) -> str
"""
Configure the network used for analysis.
:param name: name of the network to set. If `None`, a name will be generated using prefix.
:type name: string
:param prefix: prefix to prepend to auto-generated network names if name is empty
:type name: string
:return: The name of the configured network, if configured successfully.
:rtype: string
:raises BatfishException: if configuration fails
"""
if name is None:
name = prefix + get_uuid()
validate_name(name, "network")
try:
net = restv2helper.get_network(bf_session, name)
bf_session.network = str(net['name'])
return bf_session.network
except HTTPError as e:
if e.response.status_code != 404:
raise BatfishException('Unknown error accessing network', e)
json_data = workhelper.get_data_init_network(bf_session, name)
json_response = resthelper.get_json_response(
bf_session, CoordConsts.SVC_RSC_INIT_NETWORK, json_data)
network_name = json_response.get(CoordConsts.SVC_KEY_NETWORK_NAME)
if network_name is None:
raise BatfishException(
"Network initialization failed. Server response: {}".format(
json_response))
bf_session.network = str(network_name)
return bf_session.network
def bf_init_snapshot(upload, name=None, overwrite=False, background=False):
# type: (str, Optional[str], bool, bool) -> Union[str, Dict[str, str]]
"""Initialize a new snapshot.
:param upload: snapshot to upload
:type upload: zip file or directory
:param name: name of the snapshot to initialize
:type name: string
:param overwrite: whether or not to overwrite an existing snapshot with the
same name
:type overwrite: bool
:param background: whether or not to run the task in the background
:type background: bool
:return: name of initialized snapshot, or JSON dictionary of task status if background=True
:rtype: Union[str, Dict]
"""
if bf_session.network is None:
bf_set_network()
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in bf_list_snapshots():
if overwrite:
bf_delete_snapshot(name)
else:
raise ValueError('A snapshot named '
'{}'
' already exists in network '
'{}'
''.format(name, bf_session.network))
file_to_send = upload
if os.path.isdir(upload):
temp_zip_file = tempfile.NamedTemporaryFile()
zip_dir(upload, temp_zip_file)
file_to_send = temp_zip_file.name
json_data = workhelper.get_data_upload_snapshot(bf_session, name,
file_to_send)
resthelper.get_json_response(bf_session,
CoordConsts.SVC_RSC_UPLOAD_SNAPSHOT,
json_data)
work_item = workhelper.get_workitem_parse(bf_session, name)
answer_dict = workhelper.execute(work_item,
bf_session,
background=background)
if background:
bf_session.baseSnapshot = name
return answer_dict
status = WorkStatusCode(answer_dict["status"])
if status != WorkStatusCode.TERMINATEDNORMALLY:
raise BatfishException(
'Initializing snapshot {ss} failed with status {status}'.format(
ss=name, status=status))
else:
bf_session.baseSnapshot = name
bf_logger.info("Default snapshot is now set to %s",
bf_session.baseSnapshot)
return bf_session.baseSnapshot
def bf_fork_snapshot(base_name,
name=None,
overwrite=False,
background=False,
deactivate_interfaces=None,
deactivate_links=None,
deactivate_nodes=None,
restore_interfaces=None,
restore_links=None,
restore_nodes=None,
add_files=None):
# type: (str, Optional[str], bool, bool, Optional[List[Interface]], Optional[List[Edge]], Optional[List[str]], Optional[List[Interface]], Optional[List[Edge]], Optional[List[str]], Optional[str]) -> Union[str, Dict, None]
"""Copy an existing snapshot and deactivate or reactivate specified interfaces, nodes, and links on the copy.
:param base_name: name of the snapshot to copy
:type base_name: string
:param name: name of the snapshot to initialize
:type name: string
:param overwrite: whether or not to overwrite an existing snapshot with the
same name
:type overwrite: bool
:param background: whether or not to run the task in the background
:type background: bool
:param deactivate_interfaces: list of interfaces to deactivate in new snapshot
:type deactivate_interfaces: list[Interface]
:param deactivate_links: list of links to deactivate in new snapshot
:type deactivate_links: list[Edge]
:param deactivate_nodes: list of names of nodes to deactivate in new snapshot
:type deactivate_nodes: list[str]
:param restore_interfaces: list of interfaces to reactivate
:type restore_interfaces: list[Interface]
:param restore_links: list of links to reactivate
:type restore_links: list[Edge]
:param restore_nodes: list of names of nodes to reactivate
:type restore_nodes: list[str]
:param add_files: path to zip file or directory containing files to add
:type add_files: str
:return: name of initialized snapshot, JSON dictionary of task status if
background=True, or None if the call fails
:rtype: Union[str, Dict, None]
"""
if bf_session.network is None:
raise ValueError('Network must be set to fork a snapshot.')
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in bf_list_snapshots():
if overwrite:
bf_delete_snapshot(name)
else:
raise ValueError('A snapshot named '
'{}'
' already exists in network '
'{}'
''.format(name, bf_session.network))
encoded_file = None
if add_files is not None:
file_to_send = add_files
if os.path.isdir(add_files):
temp_zip_file = tempfile.NamedTemporaryFile()
zip_dir(add_files, temp_zip_file)
file_to_send = temp_zip_file.name
if os.path.isfile(file_to_send):
with open(file_to_send, "rb") as f:
encoded_file = base64.b64encode(f.read()).decode('ascii')
json_data = {
"snapshotBase": base_name,
"snapshotNew": name,
"deactivateInterfaces": deactivate_interfaces,
"deactivateLinks": deactivate_links,
"deactivateNodes": deactivate_nodes,
"restoreInterfaces": restore_interfaces,
"restoreLinks": restore_links,
"restoreNodes": restore_nodes,
"zipFile": encoded_file
}
restv2helper.fork_snapshot(bf_session, json_data)
work_item = workhelper.get_workitem_parse(bf_session, name)
answer_dict = workhelper.execute(work_item,
bf_session,
background=background)
if background:
bf_session.baseSnapshot = name
return answer_dict
status = WorkStatusCode(answer_dict['status'])
if status != WorkStatusCode.TERMINATEDNORMALLY:
raise BatfishException(
'Forking snapshot {ss} from {base} failed with status {status}'.
format(ss=name, base=base_name, status=status))
else:
bf_session.baseSnapshot = name
bf_logger.info("Default snapshot is now set to %s",
bf_session.baseSnapshot)
return bf_session.baseSnapshot
def bf_fork_snapshot(base_name,
name=None,
overwrite=False,
background=False,
deactivate_interfaces=None,
deactivate_links=None,
deactivate_nodes=None,
restore_interfaces=None,
restore_links=None,
restore_nodes=None,
add_files=None,
extra_args=None):
# type: (str, Optional[str], bool, bool, Optional[List[Interface]], Optional[List[Edge]], Optional[List[str]], Optional[List[Interface]], Optional[List[Edge]], Optional[List[str]], Optional[str], Optional[Dict[str, Any]]) -> Union[str, Dict, None]
"""Copy an existing snapshot and deactivate or reactivate specified interfaces, nodes, and links on the copy.
:param base_name: name of the snapshot to copy
:type base_name: string
:param name: name of the snapshot to initialize
:type name: string
:param overwrite: whether or not to overwrite an existing snapshot with the
same name
:type overwrite: bool
:param background: whether or not to run the task in the background
:type background: bool
:param deactivate_interfaces: list of interfaces to deactivate in new snapshot
:type deactivate_interfaces: list[Interface]
:param deactivate_links: list of links to deactivate in new snapshot
:type deactivate_links: list[Edge]
:param deactivate_nodes: list of names of nodes to deactivate in new snapshot
:type deactivate_nodes: list[str]
:param restore_interfaces: list of interfaces to reactivate
:type restore_interfaces: list[Interface]
:param restore_links: list of links to reactivate
:type restore_links: list[Edge]
:param restore_nodes: list of names of nodes to reactivate
:type restore_nodes: list[str]
:param add_files: path to zip file or directory containing files to add
:type add_files: str
:param extra_args: extra arguments to be passed to the parse command. See bf_session.additionalArgs.
:type extra_args: dict
:return: name of initialized snapshot, JSON dictionary of task status if
background=True, or None if the call fails
:rtype: Union[str, Dict, None]
"""
if bf_session.network is None:
raise ValueError('Network must be set to fork a snapshot.')
if name is None:
name = Options.default_snapshot_prefix + get_uuid()
validate_name(name)
if name in bf_list_snapshots():
if overwrite:
bf_delete_snapshot(name)
else:
raise ValueError('A snapshot named '
'{}'
' already exists in network '
'{}'
''.format(name, bf_session.network))
encoded_file = None
if add_files is not None:
file_to_send = add_files
if os.path.isdir(add_files):
temp_zip_file = tempfile.NamedTemporaryFile()
zip_dir(add_files, temp_zip_file)
file_to_send = temp_zip_file.name
if os.path.isfile(file_to_send):
with open(file_to_send, "rb") as f:
encoded_file = base64.b64encode(f.read()).decode('ascii')
json_data = {
"snapshotBase": base_name,
"snapshotNew": name,
"deactivateInterfaces": deactivate_interfaces,
"deactivateLinks": deactivate_links,
"deactivateNodes": deactivate_nodes,
"restoreInterfaces": restore_interfaces,
"restoreLinks": restore_links,
"restoreNodes": restore_nodes,
"zipFile": encoded_file
}
restv2helper.fork_snapshot(bf_session, json_data)
return _parse_snapshot(name, background, extra_args)
def test_validate_question_name():
assert validate_name('goodname')
for name in [42, 'x' * 151, '/', '/etc']:
with pytest.raises(QuestionValidationException):
assert validate_question_name(name)
def test_validate_question_name():
assert validate_name("goodname")
for name in [42, "x" * 151, "/", "/etc"]:
with pytest.raises(QuestionValidationException):
assert validate_question_name(name)
