<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  king_phisher/client/export.py
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above
#    copyright notice, this list of conditions and the following disclaimer
#    in the documentation and/or other materials provided with the
#    distribution.
#  * Neither the name of the project nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
#  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#  OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import collections
import copy
import csv
import datetime
import io
import ipaddress
import json
import logging
import os
import re
import shutil
import tarfile
import xml.etree.ElementTree as ET

from king_phisher.errors import KingPhisherInputValidationError

from boltons import iterutils
import geojson
from smoke_zephyr.utilities import escape_single_quote
from smoke_zephyr.utilities import unescape_single_quote

__all__ = [
	'campaign_to_xml',
	'convert_value',
	'message_data_to_kpm',
	'treeview_liststore_to_csv'
]

KPM_ARCHIVE_FILES = {
	'attachment_file': 'message_attachment.bin',
	'target_file': 'target_file.csv'
}

KPM_INLINE_IMAGE_REGEXP = re.compile(r"""{{\s*inline_image\(\s*(('(?:[^'\\]|\\.)+')|("(?:[^"\\]|\\.)+"))\s*\)\s*}}""")

logger = logging.getLogger('KingPhisher.Client.export')

def message_template_to_kpm(template):
	files = []
	cursor = 0
	match = True
	while match:
		match = KPM_INLINE_IMAGE_REGEXP.search(template[cursor:])
		if not match:
			break
		file_path = unescape_single_quote(match.group(1)[1:-1])
		files.append(file_path)
		file_name = os.path.basename(file_path)
		start = cursor + match.start()
		end = cursor + match.end()
		inline_tag = "{{{{ inline_image('{0}') }}}}".format(escape_single_quote(file_name))
		template = template[:start] + inline_tag + template[end:]
		cursor = start + len(inline_tag)
	return template, files

def message_template_from_kpm(template, files):
	files = dict(zip(map(os.path.basename, files), files))
	cursor = 0
	match = True
	while match:
		match = KPM_INLINE_IMAGE_REGEXP.search(template[cursor:])
		if not match:
			break
		file_name = unescape_single_quote(match.group(1)[1:-1])
		file_path = files.get(file_name)
		start = cursor + match.start()
		end = cursor + match.end()
		if not file_path:
			cursor = end
			continue
		insert_tag = "{{{{ inline_image('{0}') }}}}".format(escape_single_quote(file_path))
		template = template[:start] + insert_tag + template[end:]
		cursor = start + len(insert_tag)
	return template

def convert_value(table_name, key, value):
	"""
	Perform any conversions necessary to neatly display the data in XML
	format.

	:param str table_name: The table name that the key and value pair are from.
	:param str key: The data key.
	:param value: The data value to convert.
	:return: The converted value.
	:rtype: str
	"""
	if isinstance(value, datetime.datetime):
		value = value.isoformat()
	if value != None:
		value = str(value)
	return value

def campaign_to_xml(rpc, campaign_id, xml_file):
	"""
	Load all information for a particular campaign and dump it to an XML
	file.

	:param rpc: The connected RPC instance to load the information with.
	:type rpc: :py:class:`.KingPhisherRPCClient`
	:param campaign_id: The ID of the campaign to load the information for.
	:param str xml_file: The destination file for the XML data.
	"""
	root = ET.Element('king_phisher')
	# Generate export metadata
	metadata = ET.SubElement(root, 'metadata')
	ET.SubElement(metadata, 'timestamp').text = datetime.datetime.utcnow().isoformat()
	ET.SubElement(metadata, 'utctime').text = 'True'
	ET.SubElement(metadata, 'version').text = '1.2'

	campaign = ET.SubElement(root, 'campaign')
	campaign_info = rpc.remote_table_row('campaigns', campaign_id)
	for key, value in campaign_info._asdict().items():
		ET.SubElement(campaign, key).text = convert_value('campaigns', key, value)

	# Tables with a campaign_id field
	for table_name in ['landing_pages', 'messages', 'visits', 'credentials', 'deaddrop_deployments', 'deaddrop_connections']:
		table_element = ET.SubElement(campaign, table_name)
		for table_row in rpc.remote_table('campaign/' + table_name, campaign_id):
			table_row_element = ET.SubElement(table_element, table_name[:-1])
			for key, value in table_row._asdict().items():
				ET.SubElement(table_row_element, key).text = convert_value(table_name, key, value)

	element_tree = ET.ElementTree(root)
	element_tree.write(xml_file, encoding='utf-8', xml_declaration=True)

def campaign_visits_to_geojson(rpc, campaign_id, geojson_file):
	"""
	Export the geo location information for all the visits of a campaign into
	the `GeoJSON <http://geojson.org/>`_ format.

	:param rpc: The connected RPC instance to load the information with.
	:type rpc: :py:class:`.KingPhisherRPCClient`
	:param campaign_id: The ID of the campaign to load the information for.
	:param str geojson_file: The destination file for the GeoJSON data.
	"""
	ips_for_georesolution = {}
	ip_counter = collections.Counter()
	for visit in rpc.remote_table('campaign/visits', campaign_id):
		ip_counter.update((visit.visitor_ip,))
		visitor_ip = ipaddress.ip_address(visit.visitor_ip)
		if not isinstance(visitor_ip, ipaddress.IPv4Address):
			continue
		if visitor_ip.is_loopback or visitor_ip.is_private:
			continue
		if not visitor_ip in ips_for_georesolution:
			ips_for_georesolution[visitor_ip] = visit.first_visit
		elif ips_for_georesolution[visitor_ip] > visit.first_visit:
			ips_for_georesolution[visitor_ip] = visit.first_visit
	ips_for_georesolution = [ip for (ip, _) in sorted(ips_for_georesolution.items(), key=lambda x: x[1])]
	locations = {}
	for ip_addresses in iterutils.chunked(ips_for_georesolution, 50):
		locations.update(rpc.geoip_lookup_multi(ip_addresses))
	points = []
	for ip, location in locations.items():
		if not (location.coordinates and location.coordinates[0] and location.coordinates[1]):
			continue
		points.append(geojson.Feature(geometry=location, properties=dict(count=ip_counter[ip], ip_address=ip)))
	feature_collection = geojson.FeatureCollection(points)
	with open(geojson_file, 'w') as file_h:
		json.dump(feature_collection, file_h, indent=2, separators=(',', ': '))

def message_data_from_kpm(target_file, dest_dir):
	"""
	Retrieve the stored details describing a message from a previously exported
	file.

	:param str target_file: The file to load as a message archive.
	:param str dest_dir: The directory to extract data and attachment files to.
	:return: The restored details from the message config.
	:rtype: dict
	"""
	if not tarfile.is_tarfile(target_file):
		logger.warning('the file is not recognized as a valid tar archive')<commit_msg>Use dashes instead of underlines for GeoJSON properties<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  king_phisher/client/export.py
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above
#    copyright notice, this list of conditions and the following disclaimer
#    in the documentation and/or other materials provided with the
#    distribution.
#  * Neither the name of the project nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
#  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#  OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import collections
import copy
import csv
import datetime
import io
import ipaddress
import json
import logging
import os
import re
import shutil
import tarfile
import xml.etree.ElementTree as ET

from king_phisher.errors import KingPhisherInputValidationError

from boltons import iterutils
import geojson
from smoke_zephyr.utilities import escape_single_quote
from smoke_zephyr.utilities import unescape_single_quote

__all__ = [
	'campaign_to_xml',
	'convert_value',
	'message_data_to_kpm',
	'treeview_liststore_to_csv'
]

KPM_ARCHIVE_FILES = {
	'attachment_file': 'message_attachment.bin',
	'target_file': 'target_file.csv'
}

KPM_INLINE_IMAGE_REGEXP = re.compile(r"""{{\s*inline_image\(\s*(('(?:[^'\\]|\\.)+')|("(?:[^"\\]|\\.)+"))\s*\)\s*}}""")

logger = logging.getLogger('KingPhisher.Client.export')

def message_template_to_kpm(template):
	files = []
	cursor = 0
	match = True
	while match:
		match = KPM_INLINE_IMAGE_REGEXP.search(template[cursor:])
		if not match:
			break
		file_path = unescape_single_quote(match.group(1)[1:-1])
		files.append(file_path)
		file_name = os.path.basename(file_path)
		start = cursor + match.start()
		end = cursor + match.end()
		inline_tag = "{{{{ inline_image('{0}') }}}}".format(escape_single_quote(file_name))
		template = template[:start] + inline_tag + template[end:]
		cursor = start + len(inline_tag)
	return template, files

def message_template_from_kpm(template, files):
	files = dict(zip(map(os.path.basename, files), files))
	cursor = 0
	match = True
	while match:
		match = KPM_INLINE_IMAGE_REGEXP.search(template[cursor:])
		if not match:
			break
		file_name = unescape_single_quote(match.group(1)[1:-1])
		file_path = files.get(file_name)
		start = cursor + match.start()
		end = cursor + match.end()
		if not file_path:
			cursor = end
			continue
		insert_tag = "{{{{ inline_image('{0}') }}}}".format(escape_single_quote(file_path))
		template = template[:start] + insert_tag + template[end:]
		cursor = start + len(insert_tag)
	return template

def convert_value(table_name, key, value):
	"""
	Perform any conversions necessary to neatly display the data in XML
	format.

	:param str table_name: The table name that the key and value pair are from.
	:param str key: The data key.
	:param value: The data value to convert.
	:return: The converted value.
	:rtype: str
	"""
	if isinstance(value, datetime.datetime):
		value = value.isoformat()
	if value != None:
		value = str(value)
	return value

def campaign_to_xml(rpc, campaign_id, xml_file):
	"""
	Load all information for a particular campaign and dump it to an XML
	file.

	:param rpc: The connected RPC instance to load the information with.
	:type rpc: :py:class:`.KingPhisherRPCClient`
	:param campaign_id: The ID of the campaign to load the information for.
	:param str xml_file: The destination file for the XML data.
	"""
	root = ET.Element('king_phisher')
	# Generate export metadata
	metadata = ET.SubElement(root, 'metadata')
	ET.SubElement(metadata, 'timestamp').text = datetime.datetime.utcnow().isoformat()
	ET.SubElement(metadata, 'utctime').text = 'True'
	ET.SubElement(metadata, 'version').text = '1.2'

	campaign = ET.SubElement(root, 'campaign')
	campaign_info = rpc.remote_table_row('campaigns', campaign_id)
	for key, value in campaign_info._asdict().items():
		ET.SubElement(campaign, key).text = convert_value('campaigns', key, value)

	# Tables with a campaign_id field
	for table_name in ['landing_pages', 'messages', 'visits', 'credentials', 'deaddrop_deployments', 'deaddrop_connections']:
		table_element = ET.SubElement(campaign, table_name)
		for table_row in rpc.remote_table('campaign/' + table_name, campaign_id):
			table_row_element = ET.SubElement(table_element, table_name[:-1])
			for key, value in table_row._asdict().items():
				ET.SubElement(table_row_element, key).text = convert_value(table_name, key, value)

	element_tree = ET.ElementTree(root)
	element_tree.write(xml_file, encoding='utf-8', xml_declaration=True)

def campaign_visits_to_geojson(rpc, campaign_id, geojson_file):
	"""
	Export the geo location information for all the visits of a campaign into
	the `GeoJSON <http://geojson.org/>`_ format.

	:param rpc: The connected RPC instance to load the information with.
	:type rpc: :py:class:`.KingPhisherRPCClient`
	:param campaign_id: The ID of the campaign to load the information for.
	:param str geojson_file: The destination file for the GeoJSON data.
	"""
	ips_for_georesolution = {}
	ip_counter = collections.Counter()
	for visit in rpc.remote_table('campaign/visits', campaign_id):
		ip_counter.update((visit.visitor_ip,))
		visitor_ip = ipaddress.ip_address(visit.visitor_ip)
		if not isinstance(visitor_ip, ipaddress.IPv4Address):
			continue
		if visitor_ip.is_loopback or visitor_ip.is_private:
			continue
		if not visitor_ip in ips_for_georesolution:
			ips_for_georesolution[visitor_ip] = visit.first_visit
		elif ips_for_georesolution[visitor_ip] > visit.first_visit:
			ips_for_georesolution[visitor_ip] = visit.first_visit
	ips_for_georesolution = [ip for (ip, _) in sorted(ips_for_georesolution.items(), key=lambda x: x[1])]
	locations = {}
	for ip_addresses in iterutils.chunked(ips_for_georesolution, 50):
		locations.update(rpc.geoip_lookup_multi(ip_addresses))
	points = []
	for ip, location in locations.items():
		if not (location.coordinates and location.coordinates[0] and location.coordinates[1]):
			continue
		points.append(geojson.Feature(geometry=location, properties={'count': ip_counter[ip], 'ip-address': ip}))
	feature_collection = geojson.FeatureCollection(points)
	with open(geojson_file, 'w') as file_h:
		json.dump(feature_collection, file_h, indent=2, separators=(',', ': '))

def message_data_from_kpm(target_file, dest_dir):
	"""
	Retrieve the stored details describing a message from a previously exported
	file.

	:param str target_file: The file to load as a message archive.
	:param str dest_dir: The directory to extract data and attachment files to.
	:return: The restored details from the message config.
	:rtype: dict
	"""
	if not tarfile.is_tarfile(target_file):
		logger.warning('the file is not recognized as a valid tar archive')
{'fraction_non_alphanumeric': 0.10360996916752313, 'fraction_numerical': 0.0028263103802672147, 'mean_word_length': 5.190457256461233, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 14}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4985212', 'n_tokens_mistral': 5153, 'n_tokens_neox': 4844, 'n_words': 1502}