<commit_before># -*- coding: iso-8859-15 -*-
import re, xpath, os, os.path, jsonpath, json
from xml.dom import minidom

# IMPORTANT: Functions for assertions *MUST* be lower case

# Status
def status(assertion, request):
	assertion.expected, assertion.received = str(assertion.value).split('|'), str(request.resp_status)
	assertion.passed = True if assertion.received in assertion.expected else False

def statusnot(assertion, request):
	assertion.expected, assertion.received = str(assertion.value).split('|'), str(request.resp_status)
	assertion.passed = True if assertion.received not in assertion.expected else False

# Headers
def headerexists(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	for k in request.resp_headers.keys():
		if k.lower() == assertion.expected.lower():
			assertion.received = k
			assertion.passed = True

def headervalueequals(assertion, request):
	if 'header' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "header" --', '-- unable to query --'
		assertion.passed = False
		return
	header = assertion.params['header']
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	for k in request.resp_headers.keys():
	    if k.lower() == header.lower():
			if request.resp_headers[k].lower() == assertion.expected.lower():
				assertion.passed = True
    			assertion.received = request.resp_headers[k]
	assertion.expected = 'Header "%s" with value "%s"' % (header, assertion.expected)

def headervaluecontains(assertion, request):
	if 'header' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "header" --', '-- unable to query --'
		assertion.passed = False
		return
	header = assertion.params['header']
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	for k in request.resp_headers.keys():
	    if k.lower() == header.lower():
			if assertion.expected.lower() in request.resp_headers[k].lower():
				assertion.passed = True
    			assertion.received = request.resp_headers[k]
	assertion.expected = 'Header "%s" with value "%s"' % (header, assertion.expected)

def headerkeyvalueequals(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	try:
	   assertion.received = request.resp_headers.get(assertion.params['name'])
	   if assertion.received == assertion.expected:
	       assertion.passed = True
	except:
	   assertion.received, assertion.passed = 'Error parsing response', False

def headerkeyvaluecsontains(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	try:
	   assertion.received = request.resp_headers[assertion.params['name']]
	   if assertion.expected in assertion.received:
	       assertion.passed = True
	except:
	   assertion.received, assertion.passed = 'Error parsing response', False

def headerkeyexists(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	try:
	   assertion.received = request.resp_headers[assertion.expected]
	   assertion.passed = True
	except:
	   assertion.received, assertion.passed = 'Error parsing response', False

# Body
def bodyany(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()[:80]
	assertion.passed = True

def bodyequals(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if assertion.expected.lower() == assertion.received.decode(encoding='ascii', errors='replace').lower() else False

def bodycontains(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if assertion.expected.lower() in assertion.received.decode(encoding='ascii', errors='replace').lower() else False

def bodycontainsoneof(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	found = False
	for expect in assertion.expected.split("|"):
		if expect.lower() in assertion.received.decode(encoding='ascii', errors='replace').lower():
			assertion.passed = True
			return
	assertion.passed = False

def bodylengthequals(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if len(assertion.received)==int(assertion.expected) else False

def bodylengthlessthan(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if len(assertion.received)<int(assertion.expected) else False

def bodylengthmorethan(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if len(assertion.received)>int(assertion.expected) else False

# Body XPath
def bodyxpathany(assertion, request):
	assertion.expected = 'Any value for query "%s"' % assertion.value
	try:
		assertion.received = xpath.findvalue(assertion.value, minidom.parseString(request.resp_body))
		assertion.passed = True if assertion.received else False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathequals(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" equals "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body))
		for expected in assertion.value.split('|'):
			if expected in assertion.received:
				assertion.passed = True
				return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathcountequals(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" has %s occurrences' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = len(xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body)))
		if str(assertion.received) in str(assertion.value).split('|'):
			assertion.passed = True
			return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathcountmorethan(assertion, request, operation=">"):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" has %s%s occurrences' % (assertion.params['query'], operation, assertion.value)
	try:
		assertion.received = len(xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body)))
		if operation==">":
			assertion.passed = True if assertion.received>int(assertion.value) else False
		else:
			assertion.passed = True if assertion.received<int(assertion.value) else False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathcountlessthan(assertion, request):
	return bodyxpathcountmorethan(assertion, request, operation="<")

def bodyxpathordered(assertion, request):
	assertion.expected, assertion.received, assertion.passed = '', '', False
	pass

def bodyxpathallcontainsoneof(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" contains "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = xpath.findvalues(arg, minidom.parseString(request.response.body))
		expectedList = assertion.value.split('|')
		passed = True
		if (len(assertion.received) > 0):
			for r in assertion.received:
				found = False
				for e in expectedList:
					if (e in r):
						found = True
				if (not found):
					passed = False
	except:
		received, passed = 'error parsing response xml', False

def bodyxpathequalscount(assertion, request):
	valuesFound = None
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	if 'value' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" with Value "%s" has %s occurrences' % (assertion.params['query'], assertion.params['value'], assertion.value)
	try:
		valuesFound = xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body))
		assertion.received = valuesFound.count(assertion.params['value'])
		if str(assertion.received) in str(assertion.value).split('|'):
			assertion.passed = True
			return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def filesizeeequals(assertion, request, operator='='):
	assertion.expected, assertion.received = assertion.value, '?'
	try:
		assertion.received = int(os.path.getsize('resources/%s' % arg))
		passed = False
		if operator=='=':
			passed = True if int(assertion.expected) == int(assertion.received) else False
		elif operator=='>':
			passed = True if int(assertion.received) > int(assertion.expected) else False
		elif operator=='<':
			passed = True if int(assertion.received) < int(assertion.expected) else False
	except:
		received, passed = 'error accessing file', False

def filesizeemorethan(assertion, request):
	return filesizeeequals(assertion, request, operator='>')

def filesizeelessthan(assertion, request):
	return filesizeeequals(assertion, request, operator='<')

# bodyxpathcontains, bodyxpathallcontains, bodyxpathallequalsoneof, bodyxpathallvaluesordered, bodyxpathallvaluesmorethan
# bodyxpathallvalueslessthan, bodyxpathallvaluesmoreorequalsthan, bodyxpathallvalueslessorequalsthan, bodyxpathallvaluesequalsthan
# bodyxpathvaluemorethan bodyxpathvaluemoreorequalsthan bodyxpathvaluelessthan

# bodyregexequals bodyregexany bodyregexcountequals bodyregexcountlessthan bodyregexcountmorethan

# timelessthan timemorethan bodymatchesschema

# Body JsonPath
def bodyjsonxpathbiggestany(assertion, request): return bodyjsonpathbiggestany(assertion, request)
def bodyjsonpathbiggestany(assertion, request):
	assertion.expected = 'Any value for biggest value in query "%s"' % assertion.value
	try:
		biggest = 0
		for val in jsonpath.jsonpath(json.loads(request.resp_body), assertion.value):
			if val > biggest: biggest = val
		assertion.received = biggest
		assertion.passed = True if assertion.received else False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathany(assertion, request): return bodyjsonpathany(assertion, request)
def bodyjsonpathany(assertion, request):
	assertion.expected = 'Any value for query "%s"' % assertion.value
	try:
		assertion.received = jsonpath.jsonpath(json.loads(request.resp_body), assertion.value)[0]
		assertion.passed = True if assertion.received else False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathequals(assertion, request): return bodyjsonpathequals(assertion, request)
def bodyjsonpathequals(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" equals "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = str(jsonpath.jsonpath(json.loads(request.resp_body), assertion.params['query'])[0])
		for assertion.value in assertion.value.split('|'):
			if assertion.value == assertion.received:
				assertion.passed = True
				return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathecontains(assertion, request): return bodyjsonpathcontains(assertion, request)
def bodyjsonxpathcontains(assertion, request): return bodyjsonpathcontains(assertion, request)
def bodyjsonpathcontains(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" contains "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = jsonpath.jsonpath(json.loads(request.resp_body), assertion.params['query'])
		for assertion.value in assertion.value.split('|'):
			if assertion.value in assertion.received:
				assertion.passed = True
				return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathcountequals(assertion, request, operation='='): return bodyjsonpathcountequals(assertion, request, operation)
def bodyjsonpathcountequals(assertion, request, operation='='):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" is %s %s occurrences' % (assertion.params['query'], operation, assertion.value)
	try:
		query = jsonpath.jsonpath(json.loads(unicode(request.resp_body)), assertion.params['query'])
		if not query: assertion.received = 0
		else: assertion.received = len(query)
		if operation == '<': assertion.passed = True if assertion.received < int(assertion.value) else False
		elif operation == '>': assertion.passed = True if assertion.received > int(assertion.value) else False
		else: assertion.passed = True if assertion.received == int(assertion.value) else False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathcountmorethan(assertion, request): return bodyjsonpathcountmorethan(assertion, request)
def bodyjsonpathcountmorethan(assertion, request):
    bodyjsonxpathcountequals(assertion, request, operation='>')

def bodyjsonxpathcountlessthan(assertion, request): return bodyjsonpathcountlessthan(assertion, request)
def bodyjsonpathcountlessthan(assertion, request):
    bodyjsonxpathcountequals(assertion, request, operation='<')

# Factory for assertions
def run_assertion(assertion, request):
	assertion_type = assertion.type.strip().lower()
	if assertion_type in globals().keys():
		globals()[assertion.type.strip().lower()](assertion, request)
		if assertion.params.get('inverse') == 'true': assertion.passed = not assertion.passed
		if len(str(assertion.received))>200: assertion.received = '%s[...]' % assertion.received[:200]
		return assertion.passed
	else:
		assertion.expected, assertion.received = str(assertion.value), '--Unkown Assertion Type--'
		assertion.passed = False
		return False
<commit_msg>Small cleanup<commit_after># -*- coding: iso-8859-15 -*-
import re, xpath, os, os.path, jsonpath, json
from xml.dom import minidom

# IMPORTANT: Functions for assertions *MUST* be lower case

##############################################################################
# Status
##############################################################################
def status(assertion, request):
	assertion.expected, assertion.received = str(assertion.value).split('|'), str(request.resp_status)
	assertion.passed = True if assertion.received in assertion.expected else False

def statusnot(assertion, request):
	assertion.expected, assertion.received = str(assertion.value).split('|'), str(request.resp_status)
	assertion.passed = True if assertion.received not in assertion.expected else False

# Headers
def headerexists(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	for k in request.resp_headers.keys():
		if k.lower() == assertion.expected.lower():
			assertion.received = k
			assertion.passed = True

def headervalueequals(assertion, request):
	if 'header' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "header" --', '-- unable to query --'
		assertion.passed = False
		return
	header = assertion.params['header']
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	for k in request.resp_headers.keys():
	    if k.lower() == header.lower():
			if request.resp_headers[k].lower() == assertion.expected.lower():
				assertion.passed = True
    			assertion.received = request.resp_headers[k]
	assertion.expected = 'Header "%s" with value "%s"' % (header, assertion.expected)

def headervaluecontains(assertion, request):
	if 'header' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "header" --', '-- unable to query --'
		assertion.passed = False
		return
	header = assertion.params['header']
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	for k in request.resp_headers.keys():
	    if k.lower() == header.lower():
			if assertion.expected.lower() in request.resp_headers[k].lower():
				assertion.passed = True
    			assertion.received = request.resp_headers[k]
	assertion.expected = 'Header "%s" with value "%s"' % (header, assertion.expected)

def headerkeyvalueequals(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	try:
	   assertion.received = request.resp_headers.get(assertion.params['name'])
	   if assertion.received == assertion.expected:
	       assertion.passed = True
	except:
	   assertion.received, assertion.passed = 'Error parsing response', False

def headerkeyvaluecsontains(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	try:
	   assertion.received = request.resp_headers[assertion.params['name']]
	   if assertion.expected in assertion.received:
	       assertion.passed = True
	except:
	   assertion.received, assertion.passed = 'Error parsing response', False

def headerkeyexists(assertion, request):
	assertion.expected, assertion.received = assertion.value, 'Not Found'
	assertion.passed = False
	try:
	   assertion.received = request.resp_headers[assertion.expected]
	   assertion.passed = True
	except:
	   assertion.received, assertion.passed = 'Error parsing response', False

##############################################################################
# Body
##############################################################################
def bodyany(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()[:80]
	assertion.passed = True

def bodyequals(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if assertion.expected.lower() == assertion.received.decode(encoding='ascii', errors='replace').lower() else False

def bodycontains(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if assertion.expected.lower() in assertion.received.decode(encoding='ascii', errors='replace').lower() else False

def bodycontainsoneof(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	for expect in assertion.expected.split("|"):
		if expect.lower() in assertion.received.decode(encoding='ascii', errors='replace').lower():
			assertion.passed = True
			return
	assertion.passed = False

def bodylengthequals(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if len(assertion.received)==int(assertion.expected) else False

def bodylengthlessthan(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if len(assertion.received)<int(assertion.expected) else False

def bodylengthmorethan(assertion, request):
	assertion.expected, assertion.received = assertion.value, request.resp_body.strip()
	assertion.passed = True if len(assertion.received)>int(assertion.expected) else False

# Body XPath
def bodyxpathany(assertion, request):
	assertion.expected = 'Any value for query "%s"' % assertion.value
	try:
		assertion.received = xpath.findvalue(assertion.value, minidom.parseString(request.resp_body))
		assertion.passed = True if assertion.received else False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathequals(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" equals "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body))
		for expected in assertion.value.split('|'):
			if expected in assertion.received:
				assertion.passed = True
				return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathcountequals(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" has %s occurrences' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = len(xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body)))
		if str(assertion.received) in str(assertion.value).split('|'):
			assertion.passed = True
			return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathcountmorethan(assertion, request, operation=">"):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" has %s%s occurrences' % (assertion.params['query'], operation, assertion.value)
	try:
		assertion.received = len(xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body)))
		if operation==">":
			assertion.passed = True if assertion.received>int(assertion.value) else False
		else:
			assertion.passed = True if assertion.received<int(assertion.value) else False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def bodyxpathcountlessthan(assertion, request):
	return bodyxpathcountmorethan(assertion, request, operation="<")

def bodyxpathordered(assertion, request):
	assertion.expected, assertion.received, assertion.passed = '', '', False
	pass

def bodyxpathallcontainsoneof(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" contains "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = xpath.findvalues(arg, minidom.parseString(request.response.body))
		expectedList = assertion.value.split('|')
		passed = True
		if (len(assertion.received) > 0):
			for r in assertion.received:
				found = False
				for e in expectedList:
					if (e in r):
						found = True
				if (not found):
					passed = False
	except:
		received, passed = 'error parsing response xml', False

def bodyxpathequalscount(assertion, request):
	valuesFound = None
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	if 'value' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" with Value "%s" has %s occurrences' % (assertion.params['query'], assertion.params['value'], assertion.value)
	try:
		valuesFound = xpath.findvalues(assertion.params['query'], minidom.parseString(request.resp_body))
		assertion.received = valuesFound.count(assertion.params['value'])
		if str(assertion.received) in str(assertion.value).split('|'):
			assertion.passed = True
			return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response XML', False

def filesizeeequals(assertion, request, operator='='):
	assertion.expected, assertion.received = assertion.value, '?'
	try:
		assertion.received = int(os.path.getsize('resources/%s' % arg))
		passed = False
		if operator=='=':
			passed = True if int(assertion.expected) == int(assertion.received) else False
		elif operator=='>':
			passed = True if int(assertion.received) > int(assertion.expected) else False
		elif operator=='<':
			passed = True if int(assertion.received) < int(assertion.expected) else False
	except:
		received, passed = 'error accessing file', False

def filesizeemorethan(assertion, request):
	return filesizeeequals(assertion, request, operator='>')

def filesizeelessthan(assertion, request):
	return filesizeeequals(assertion, request, operator='<')

# Body JsonPath
def bodyjsonxpathbiggestany(assertion, request): return bodyjsonpathbiggestany(assertion, request)
def bodyjsonpathbiggestany(assertion, request):
	assertion.expected = 'Any value for biggest value in query "%s"' % assertion.value
	try:
		biggest = 0
		for val in jsonpath.jsonpath(json.loads(request.resp_body), assertion.value):
			if val > biggest: biggest = val
		assertion.received = biggest
		assertion.passed = True if assertion.received else False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathany(assertion, request): return bodyjsonpathany(assertion, request)
def bodyjsonpathany(assertion, request):
	assertion.expected = 'Any value for query "%s"' % assertion.value
	try:
		assertion.received = jsonpath.jsonpath(json.loads(request.resp_body), assertion.value)[0]
		assertion.passed = True if assertion.received else False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathequals(assertion, request): return bodyjsonpathequals(assertion, request)
def bodyjsonpathequals(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" equals "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = str(jsonpath.jsonpath(json.loads(request.resp_body), assertion.params['query'])[0])
		for assertion.value in assertion.value.split('|'):
			if assertion.value == assertion.received:
				assertion.passed = True
				return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathecontains(assertion, request): return bodyjsonpathcontains(assertion, request)
def bodyjsonxpathcontains(assertion, request): return bodyjsonpathcontains(assertion, request)
def bodyjsonpathcontains(assertion, request):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" contains "%s"' % (assertion.params['query'], assertion.value)
	try:
		assertion.received = jsonpath.jsonpath(json.loads(request.resp_body), assertion.params['query'])
		for assertion.value in assertion.value.split('|'):
			if assertion.value in assertion.received:
				assertion.passed = True
				return
		assertion.passed = False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathcountequals(assertion, request, operation='='): return bodyjsonpathcountequals(assertion, request, operation)
def bodyjsonpathcountequals(assertion, request, operation='='):
	if 'query' not in assertion.params.keys():
		assertion.expected, assertion.received = '-- Missing assertion attribute "query" --', '-- unable to query --'
		assertion.passed = False
		return
	assertion.expected = 'Query "%s" is %s %s occurrences' % (assertion.params['query'], operation, assertion.value)
	try:
		query = jsonpath.jsonpath(json.loads(unicode(request.resp_body)), assertion.params['query'])
		if not query: assertion.received = 0
		else: assertion.received = len(query)
		if operation == '<': assertion.passed = True if assertion.received < int(assertion.value) else False
		elif operation == '>': assertion.passed = True if assertion.received > int(assertion.value) else False
		else: assertion.passed = True if assertion.received == int(assertion.value) else False
	except:
		assertion.received, assertion.passed = 'Error parsing response Json', False

def bodyjsonxpathcountmorethan(assertion, request): return bodyjsonpathcountmorethan(assertion, request)
def bodyjsonpathcountmorethan(assertion, request):
    bodyjsonxpathcountequals(assertion, request, operation='>')

def bodyjsonxpathcountlessthan(assertion, request): return bodyjsonpathcountlessthan(assertion, request)
def bodyjsonpathcountlessthan(assertion, request):
    bodyjsonxpathcountequals(assertion, request, operation='<')

# Factory for assertions
def run_assertion(assertion, request):
	assertion_type = assertion.type.strip().lower()
	if assertion_type in globals().keys():
		globals()[assertion.type.strip().lower()](assertion, request)
		if assertion.params.get('inverse') == 'true': assertion.passed = not assertion.passed
		if len(str(assertion.received))>200: assertion.received = '%s[...]' % assertion.received[:200]
		return assertion.passed
	else:
		assertion.expected, assertion.received = str(assertion.value), '--Unkown Assertion Type--'
		assertion.passed = False
		return False
<|endoftext|>


<commit_before>import boto3
import contextlib
import gc
import gzip
import os
import pandas as pd
import Queue
import sys
import threading
import time
import ujson as json
import urllib
import urllib2
import uuid

from bisect import bisect
from boto3.s3.transfer import S3Transfer
from datetime import datetime, timedelta
from moztelemetry import get_pings_properties
from moztelemetry.dataset import Dataset
from sets import Set
from StringIO import StringIO

from profile import ProfileProcessor

UNSYMBOLICATED = "<unsymbolicated>"
REDUCE_BY_KEY_PARALLELISM = 512

def time_code(name, fn):
    print "{}...".format(name)
    start = time.time()
    result = fn()
    end = time.time()
    delta = end - start
    print "{} took {}ms to complete".format(name, int(round(delta * 1000)))
    return result

def get_data(sc, config, start_date_relative, end_date_relative):
    start_date = (datetime.today() + timedelta(days=start_date_relative))
    start_date_str = start_date.strftime("%Y%m%d")
    end_date = (datetime.today() + timedelta(days=end_date_relative))
    end_date_str = end_date.strftime("%Y%m%d")

    pings = (Dataset.from_source("telemetry")
        .where(docType='OTHER')
        .where(submissionDate=lambda b: (b.startswith(start_date_str) or b > start_date_str)
                                         and (b.startswith(end_date_str) or b < end_date_str))
        .where(appUpdateChannel="nightly")
        .records(sc, sample=config['sample_size']))

    pings = pings.filter(lambda p: p.get('meta', {}).get('docType', {}) == 'bhr')

    properties = ["environment/system/os/name",
                  "meta/submissionDate",
                  "payload/modules",
                  "payload/hangs"]

    try:
        return get_pings_properties(pings, properties, with_processes=True)
    except ValueError:
        return None

def windows_only(p):
    return p["environment/system/os/name"] == "Windows_NT"

def ping_is_valid(ping):
    if not isinstance(ping["meta/submissionDate"], basestring):
        return False
    # if type(ping["payload/info/subsessionLength"]) != int:
    #     return False

    return True

def process_stack(stack, modules):
    if isinstance(stack, list):
        if stack[0] == -1 or stack[0] is None:
            return (None, stack[1])
        return ((modules[stack[0]][0], modules[stack[0]][1]), stack[1])
    else:
        return (('pseudo', None), stack)

def process_hangs(ping):
    result = []

    submission_date = ping["meta/submissionDate"][:8] # "YYYYMMDD" : 8 characters
    # usage_hours = float(ping['payload/info/subsessionLength']) / 60.0

    # if usage_hours == 0:
    #     return result
    modules = ping['payload/modules']
    hangs = ping['payload/hangs']

    return [(
        [process_stack(s, modules) for s in h['stack']],
        h['duration'],
        h['thread'],
        h['runnableName'],
        h['process'],
        h['annotations'],
        submission_date,
    ) for h in hangs]

def get_all_hangs(pings):
    return pings.flatMap(lambda p: process_hangs(p))

def map_to_frame_info(hang):
    memory_map = hang['hang']['nativeStack']['memoryMap']
    stack = hang['hang']['nativeStack']['stacks'][0]
    return [
        (tuple(memory_map[module_index]), (offset,)) if module_index != -1 else (None, (offset,))
        for module_index, offset in stack
    ]

def get_frames_by_module(hangs):
    return (hangs.flatMap(lambda hang: hang[0])
        .filter(lambda hang_tuple: hang_tuple[0] is not None)
        .map(lambda hang_tuple: (hang_tuple[0], (hang_tuple[1],)))
        .distinct()
        .reduceByKey(lambda a,b: a + b, REDUCE_BY_KEY_PARALLELISM)
        .collectAsMap())

def symbolicate_stacks(stack, processed_modules):
    symbolicated = []
    for module, offset in stack:
        if module is not None:
            debug_name, breakpad_id = module
            processed = processed_modules[breakpad_id, offset]
            if processed is not None:
                symbolicated.append(processed)
            else:
                symbolicated.append((UNSYMBOLICATED, debug_name))
        else:
            symbolicated.append((UNSYMBOLICATED, 'unknown'))
    return symbolicated

def map_to_hang_data(hang, config):
    stack, duration, thread, runnable_name, process, annotations, submission_date = hang
    if duration < config['hang_lower_bound']:
        return None
    if duration >= config['hang_upper_bound']:
        return None

    user_interacting = False
    if 'UserInteracting' in annotations and annotations['UserInteracting'] == 'true':
        user_interacting = True

    key = (
        tuple((a,b) for a,b in stack),
        runnable_name,
        thread,
        submission_date,
        user_interacting)
    return (key, (
        float(duration),
        1.0,
    ))

def merge_hang_data(a, b):
    return (
        a[0] + b[0],
        a[1] + b[1],
    )

def process_hang_key(key, processed_modules):
    stack, runnable_name, thread, submission_date, user_interacting = key
    symbolicated = symbolicate_stacks(stack, processed_modules)

    return (
        tuple(symbolicated),
        runnable_name,
        thread,
        submission_date,
        user_interacting
    )

def process_hang_value(key, val, usage_hours_by_date):
    return val

def get_grouped_sums_and_counts(hangs, processed_modules, usage_hours_by_date, config):
    reduced = (hangs
        .map(lambda hang: map_to_hang_data(hang, config))
        .filter(lambda hang: hang is not None)
        .reduceByKey(merge_hang_data, REDUCE_BY_KEY_PARALLELISM)
        .collect())
    items = [
        (process_hang_key(k, processed_modules), process_hang_value(k, v, usage_hours_by_date))
        for k, v in reduced
    ]
    return [
        k + v for k, v in items if k is not None
    ]

def get_usage_hours(ping):
    submission_date = ping["meta/submissionDate"]
    usage_hours = float(ping['payload/info/subsessionLength']) / 60.0 / 60.0
    return (submission_date, usage_hours)

def merge_usage_hours(a, b):
    return a + b

def get_usage_hours_by_date(pings):
    return (pings.map(get_usage_hours)
        .reduceByKey(merge_usage_hours, REDUCE_BY_KEY_PARALLELISM)
        .collectAsMap())

def make_sym_map(data):
    public_symbols = {}
    func_symbols = {}

    for line in data.splitlines():
        if line.startswith("PUBLIC "):
            line = line.rstrip()
            fields = line.split(" ", 3)
            if len(fields) < 4:
                continue
            address = int(fields[1], 16)
            symbol = fields[3]
            public_symbols[address] = symbol
        elif line.startswith("FUNC "):
            line = line.rstrip()
            fields = line.split(" ", 4)
            if len(fields) < 5:
                continue
            address = int(fields[1], 16)
            symbol = fields[4]
            func_symbols[address] = symbol
    # Prioritize PUBLIC symbols over FUNC ones
    sym_map = func_symbols
    sym_map.update(public_symbols)

    return sorted(sym_map), sym_map

def get_file_URL(module, config):
    lib_name, breakpad_id = module
    if lib_name.endswith(".pdb"):
        file_name = lib_name[:-4] + ".sym"
    else:
        file_name = lib_name + ".sym"

    return config['symbol_server_url'] + "/".join([
        urllib.quote_plus(lib_name),
        urllib.quote_plus(breakpad_id),
        urllib.quote_plus(file_name)
    ])

def process_module(module, offsets, config):
    result = []
    if module[0] == 'pseudo':
        return [((None, offset), (offset, '')) for offset in offsets]
    file_URL = get_file_URL(module, config)
    success, response = fetch_URL(file_URL)
    module_name, breakpad_id = module

    if success:
        sorted_keys, sym_map = make_sym_map(response)

        for offset in offsets:
            i = bisect(sorted_keys, int(offset, 16))
            key = sorted_keys[i - 1] if i else None

            symbol = sym_map.get(key)
            if symbol is not None:
                result.append(((breakpad_id, offset), (symbol, module_name)))
            else:
                result.append(((breakpad_id, offset), (UNSYMBOLICATED, module_name)))
    else:
        for offset in offsets:
            result.append(((breakpad_id, offset), (UNSYMBOLICATED, module_name)))
    return result

def process_modules(sc, frames_by_module, config):
    data = sc.parallelize(frames_by_module.iteritems())
    return data.flatMap(lambda x: process_module(x[0], x[1], config)).collectAsMap()

def transform_pings(sc, pings, config):
    windows_pings_only = time_code("Filtering to Windows pings",
        lambda: pings.filter(windows_only).filter(ping_is_valid))

    hangs = time_code("Filtering to hangs with native stacks",
        lambda: get_all_hangs(windows_pings_only))

    frames_by_module = time_code("Getting stacks by module",
        lambda: get_frames_by_module(hangs))

    processed_modules = time_code("Processing modules",
        lambda: process_modules(sc, frames_by_module, config))

    # TODO: usage_hours_by_date is a placeholder right now until we get proper usage
    # time in our telemetry
    usage_hours_by_date = None
    # usage_hours_by_date = time_code("Getting usage hours",
    #     lambda: get_usage_hours_by_date(windows_pings_only))

    result = time_code("Grouping stacks",
        lambda: get_grouped_sums_and_counts(hangs, processed_modules, usage_hours_by_date, config))
    return result

def fetch_URL(url):
    result = False, ""
    try:
        with contextlib.closing(urllib2.urlopen(url)) as response:
            responseCode = response.getcode()
            if responseCode == 404:
                return False, ""
            if responseCode != 200:
                result = False, ""
            return True, decode_response(response)
    except IOError as e:
        result = False, ""

    if not result[0]:
        try:
            with contextlib.closing(urllib2.urlopen(url)) as response:
                responseCode = response.getcode()
                if responseCode == 404:
                    return False, ""
                if responseCode != 200:
                    result = False, ""
                return True, decode_response(response)
        except IOError as e:
            result = False, ""

    return result

def decode_response(response):
    headers = response.info()
    content_encoding = headers.get("Content-Encoding", "").lower()
    if content_encoding in ("gzip", "x-gzip", "deflate"):
        with contextlib.closing(StringIO(response.read())) as data_stream:
            try:
                with gzip.GzipFile(fileobj=data_stream) as f:
                    return f.read()
            except EnvironmentError:
                data_stream.seek(0)
                return data_stream.read().decode('zlib')
    return response.read()

def smart_hex(offset):
    return hex(int(offset))

def write_file(name, stuff, config):
    end_date = datetime.today()
    end_date_str = end_date.strftime("%Y%m%d")

    filename = "./output/%s-%s.json" % (name, end_date_str)
    gzfilename = filename + '.gz'
    jsonblob = json.dumps(stuff, ensure_ascii=False)

    if not os.path.exists('./output'):
        os.makedirs('./output')
    with gzip.open(gzfilename, 'w') as f:
        f.write(jsonblob)

    if config['use_s3']:
        bucket = "telemetry-public-analysis-2"
        s3_key = "bhr/data/hang_aggregates/" + name + ".json"
        s3_uuid_key = "bhr/data/hang_aggregates/" + name + "_" + config['uuid'] + ".json"
        client = boto3.client('s3', 'us-west-2')
        transfer = S3Transfer(client)
        transfer.upload_file(gzfilename,
                             bucket,
                             s3_key,
                             extra_args={
                                'ContentType':'application/json',
                                'ContentEncoding':'gzip'
                            })
        transfer.upload_file(gzfilename,
                             bucket,
                             s3_uuid_key,
                             extra_args={
                                'ContentType':'application/json',
                                'ContentEncoding':'gzip'
                            })

def etl_job(sc, sqlContext, config=None):
    """This is the function that will be executed on the cluster"""

    final_config = {
        'days_to_aggregate': 21,
        'date_clumping': 2,
        'use_s3': True,
        'sample_size': 0.05,
        'symbol_server_url': "https://s3-us-west-2.amazonaws.com/org.mozilla.crash-stats.symbols-public/v1/",
        'hang_profile_filename': 'hang_profile_128_16000',
        'print_debug_info': False,
        'hang_lower_bound': 128,
        'hang_upper_bound': 16000,
        'stack_acceptance_threshold': 0.001,
        'hang_outlier_threshold': 512,
        'uuid': uuid.uuid4().hex,
    }

    if config is not None:
        final_config.update(config)

    iterations = (final_config['days_to_aggregate'] + final_config['date_clumping'] - 1) / final_config['date_clumping']
    profile_processor = ProfileProcessor(final_config)

    job_start = time.time()
    # We were OOMing trying to allocate a contiguous array for all of this. Pass it in
    # bit by bit to the profile processor and hope it can handle it.
    for x in xrange(0, iterations):
        day_start = time.time()
        rel_end = x * final_config['date_clumping']
        rel_end = -rel_end - 1
        rel_start = min((x + 1) * final_config['date_clumping'] - 1, final_config['days_to_aggregate'])
        rel_start = -rel_start - 1
        data = time_code("Getting data",
            lambda: get_data(sc, final_config, rel_start, rel_end))
        if data is None:
            continue
        transformed = transform_pings(sc, data, final_config)
        time_code("Passing stacks to processor", lambda: profile_processor.ingest(transformed))
        # Run a collection to ensure that any references to any RDDs are cleaned up,
        # allowing the JVM to clean them up on its end.
        gc.collect()
        day_end = time.time()
        day_delta = day_end - day_start
        print "Iteration took {}s".format(int(round(day_delta)))
        job_elapsed = time.time() - job_start
        percent_done = float(x + 1) / float(iterations)
        projected = job_elapsed / percent_done
        remaining = projected - job_elapsed
        print "Job should finish in {}".format(timedelta(seconds=remaining))

    profile = profile_processor.process_into_profile()
    write_file(final_config['hang_profile_filename'], profile, final_config)
<commit_msg>Change UserInteracting to PendingInput<commit_after>import boto3
import contextlib
import gc
import gzip
import os
import pandas as pd
import Queue
import sys
import threading
import time
import ujson as json
import urllib
import urllib2
import uuid

from bisect import bisect
from boto3.s3.transfer import S3Transfer
from datetime import datetime, timedelta
from moztelemetry import get_pings_properties
from moztelemetry.dataset import Dataset
from sets import Set
from StringIO import StringIO

from profile import ProfileProcessor

UNSYMBOLICATED = "<unsymbolicated>"
REDUCE_BY_KEY_PARALLELISM = 512

def time_code(name, fn):
    print "{}...".format(name)
    start = time.time()
    result = fn()
    end = time.time()
    delta = end - start
    print "{} took {}ms to complete".format(name, int(round(delta * 1000)))
    return result

def get_data(sc, config, start_date_relative, end_date_relative):
    start_date = (datetime.today() + timedelta(days=start_date_relative))
    start_date_str = start_date.strftime("%Y%m%d")
    end_date = (datetime.today() + timedelta(days=end_date_relative))
    end_date_str = end_date.strftime("%Y%m%d")

    pings = (Dataset.from_source("telemetry")
        .where(docType='OTHER')
        .where(submissionDate=lambda b: (b.startswith(start_date_str) or b > start_date_str)
                                         and (b.startswith(end_date_str) or b < end_date_str))
        .where(appUpdateChannel="nightly")
        .records(sc, sample=config['sample_size']))

    pings = pings.filter(lambda p: p.get('meta', {}).get('docType', {}) == 'bhr')

    properties = ["environment/system/os/name",
                  "meta/submissionDate",
                  "payload/modules",
                  "payload/hangs"]

    try:
        return get_pings_properties(pings, properties, with_processes=True)
    except ValueError:
        return None

def windows_only(p):
    return p["environment/system/os/name"] == "Windows_NT"

def ping_is_valid(ping):
    if not isinstance(ping["meta/submissionDate"], basestring):
        return False
    # if type(ping["payload/info/subsessionLength"]) != int:
    #     return False

    return True

def process_stack(stack, modules):
    if isinstance(stack, list):
        if stack[0] == -1 or stack[0] is None:
            return (None, stack[1])
        return ((modules[stack[0]][0], modules[stack[0]][1]), stack[1])
    else:
        return (('pseudo', None), stack)

def process_hangs(ping):
    result = []

    submission_date = ping["meta/submissionDate"][:8] # "YYYYMMDD" : 8 characters
    # usage_hours = float(ping['payload/info/subsessionLength']) / 60.0

    # if usage_hours == 0:
    #     return result
    modules = ping['payload/modules']
    hangs = ping['payload/hangs']

    return [(
        [process_stack(s, modules) for s in h['stack']],
        h['duration'],
        h['thread'],
        h['runnableName'],
        h['process'],
        h['annotations'],
        submission_date,
    ) for h in hangs]

def get_all_hangs(pings):
    return pings.flatMap(lambda p: process_hangs(p))

def map_to_frame_info(hang):
    memory_map = hang['hang']['nativeStack']['memoryMap']
    stack = hang['hang']['nativeStack']['stacks'][0]
    return [
        (tuple(memory_map[module_index]), (offset,)) if module_index != -1 else (None, (offset,))
        for module_index, offset in stack
    ]

def get_frames_by_module(hangs):
    return (hangs.flatMap(lambda hang: hang[0])
        .filter(lambda hang_tuple: hang_tuple[0] is not None)
        .map(lambda hang_tuple: (hang_tuple[0], (hang_tuple[1],)))
        .distinct()
        .reduceByKey(lambda a,b: a + b, REDUCE_BY_KEY_PARALLELISM)
        .collectAsMap())

def symbolicate_stacks(stack, processed_modules):
    symbolicated = []
    for module, offset in stack:
        if module is not None:
            debug_name, breakpad_id = module
            processed = processed_modules[breakpad_id, offset]
            if processed is not None:
                symbolicated.append(processed)
            else:
                symbolicated.append((UNSYMBOLICATED, debug_name))
        else:
            symbolicated.append((UNSYMBOLICATED, 'unknown'))
    return symbolicated

def map_to_hang_data(hang, config):
    stack, duration, thread, runnable_name, process, annotations, submission_date = hang
    if duration < config['hang_lower_bound']:
        return None
    if duration >= config['hang_upper_bound']:
        return None

    pending_input = False
    if 'PendingInput' in annotations:
        pending_input = True

    key = (
        tuple((a,b) for a,b in stack),
        runnable_name,
        thread,
        submission_date,
        pending_input)
    return (key, (
        float(duration),
        1.0,
    ))

def merge_hang_data(a, b):
    return (
        a[0] + b[0],
        a[1] + b[1],
    )

def process_hang_key(key, processed_modules):
    stack, runnable_name, thread, submission_date, pending_input = key
    symbolicated = symbolicate_stacks(stack, processed_modules)

    return (
        tuple(symbolicated),
        runnable_name,
        thread,
        submission_date,
        pending_input
    )

def process_hang_value(key, val, usage_hours_by_date):
    return val

def get_grouped_sums_and_counts(hangs, processed_modules, usage_hours_by_date, config):
    reduced = (hangs
        .map(lambda hang: map_to_hang_data(hang, config))
        .filter(lambda hang: hang is not None)
        .reduceByKey(merge_hang_data, REDUCE_BY_KEY_PARALLELISM)
        .collect())
    items = [
        (process_hang_key(k, processed_modules), process_hang_value(k, v, usage_hours_by_date))
        for k, v in reduced
    ]
    return [
        k + v for k, v in items if k is not None
    ]

def get_usage_hours(ping):
    submission_date = ping["meta/submissionDate"]
    usage_hours = float(ping['payload/info/subsessionLength']) / 60.0 / 60.0
    return (submission_date, usage_hours)

def merge_usage_hours(a, b):
    return a + b

def get_usage_hours_by_date(pings):
    return (pings.map(get_usage_hours)
        .reduceByKey(merge_usage_hours, REDUCE_BY_KEY_PARALLELISM)
        .collectAsMap())

def make_sym_map(data):
    public_symbols = {}
    func_symbols = {}

    for line in data.splitlines():
        if line.startswith("PUBLIC "):
            line = line.rstrip()
            fields = line.split(" ", 3)
            if len(fields) < 4:
                continue
            address = int(fields[1], 16)
            symbol = fields[3]
            public_symbols[address] = symbol
        elif line.startswith("FUNC "):
            line = line.rstrip()
            fields = line.split(" ", 4)
            if len(fields) < 5:
                continue
            address = int(fields[1], 16)
            symbol = fields[4]
            func_symbols[address] = symbol
    # Prioritize PUBLIC symbols over FUNC ones
    sym_map = func_symbols
    sym_map.update(public_symbols)

    return sorted(sym_map), sym_map

def get_file_URL(module, config):
    lib_name, breakpad_id = module
    if lib_name.endswith(".pdb"):
        file_name = lib_name[:-4] + ".sym"
    else:
        file_name = lib_name + ".sym"

    return config['symbol_server_url'] + "/".join([
        urllib.quote_plus(lib_name),
        urllib.quote_plus(breakpad_id),
        urllib.quote_plus(file_name)
    ])

def process_module(module, offsets, config):
    result = []
    if module[0] == 'pseudo':
        return [((None, offset), (offset, '')) for offset in offsets]
    file_URL = get_file_URL(module, config)
    success, response = fetch_URL(file_URL)
    module_name, breakpad_id = module

    if success:
        sorted_keys, sym_map = make_sym_map(response)

        for offset in offsets:
            i = bisect(sorted_keys, int(offset, 16))
            key = sorted_keys[i - 1] if i else None

            symbol = sym_map.get(key)
            if symbol is not None:
                result.append(((breakpad_id, offset), (symbol, module_name)))
            else:
                result.append(((breakpad_id, offset), (UNSYMBOLICATED, module_name)))
    else:
        for offset in offsets:
            result.append(((breakpad_id, offset), (UNSYMBOLICATED, module_name)))
    return result

def process_modules(sc, frames_by_module, config):
    data = sc.parallelize(frames_by_module.iteritems())
    return data.flatMap(lambda x: process_module(x[0], x[1], config)).collectAsMap()

def transform_pings(sc, pings, config):
    windows_pings_only = time_code("Filtering to Windows pings",
        lambda: pings.filter(windows_only).filter(ping_is_valid))

    hangs = time_code("Filtering to hangs with native stacks",
        lambda: get_all_hangs(windows_pings_only))

    frames_by_module = time_code("Getting stacks by module",
        lambda: get_frames_by_module(hangs))

    processed_modules = time_code("Processing modules",
        lambda: process_modules(sc, frames_by_module, config))

    # TODO: usage_hours_by_date is a placeholder right now until we get proper usage
    # time in our telemetry
    usage_hours_by_date = None
    # usage_hours_by_date = time_code("Getting usage hours",
    #     lambda: get_usage_hours_by_date(windows_pings_only))

    result = time_code("Grouping stacks",
        lambda: get_grouped_sums_and_counts(hangs, processed_modules, usage_hours_by_date, config))
    return result

def fetch_URL(url):
    result = False, ""
    try:
        with contextlib.closing(urllib2.urlopen(url)) as response:
            responseCode = response.getcode()
            if responseCode == 404:
                return False, ""
            if responseCode != 200:
                result = False, ""
            return True, decode_response(response)
    except IOError as e:
        result = False, ""

    if not result[0]:
        try:
            with contextlib.closing(urllib2.urlopen(url)) as response:
                responseCode = response.getcode()
                if responseCode == 404:
                    return False, ""
                if responseCode != 200:
                    result = False, ""
                return True, decode_response(response)
        except IOError as e:
            result = False, ""

    return result

def decode_response(response):
    headers = response.info()
    content_encoding = headers.get("Content-Encoding", "").lower()
    if content_encoding in ("gzip", "x-gzip", "deflate"):
        with contextlib.closing(StringIO(response.read())) as data_stream:
            try:
                with gzip.GzipFile(fileobj=data_stream) as f:
                    return f.read()
            except EnvironmentError:
                data_stream.seek(0)
                return data_stream.read().decode('zlib')
    return response.read()

def smart_hex(offset):
    return hex(int(offset))

def write_file(name, stuff, config):
    end_date = datetime.today()
    end_date_str = end_date.strftime("%Y%m%d")

    filename = "./output/%s-%s.json" % (name, end_date_str)
    gzfilename = filename + '.gz'
    jsonblob = json.dumps(stuff, ensure_ascii=False)

    if not os.path.exists('./output'):
        os.makedirs('./output')
    with gzip.open(gzfilename, 'w') as f:
        f.write(jsonblob)

    if config['use_s3']:
        bucket = "telemetry-public-analysis-2"
        s3_key = "bhr/data/hang_aggregates/" + name + ".json"
        s3_uuid_key = "bhr/data/hang_aggregates/" + name + "_" + config['uuid'] + ".json"
        client = boto3.client('s3', 'us-west-2')
        transfer = S3Transfer(client)
        transfer.upload_file(gzfilename,
                             bucket,
                             s3_key,
                             extra_args={
                                'ContentType':'application/json',
                                'ContentEncoding':'gzip'
                            })
        transfer.upload_file(gzfilename,
                             bucket,
                             s3_uuid_key,
                             extra_args={
                                'ContentType':'application/json',
                                'ContentEncoding':'gzip'
                            })

def etl_job(sc, sqlContext, config=None):
    """This is the function that will be executed on the cluster"""

    final_config = {
        'days_to_aggregate': 21,
        'date_clumping': 2,
        'use_s3': True,
        'sample_size': 0.05,
        'symbol_server_url': "https://s3-us-west-2.amazonaws.com/org.mozilla.crash-stats.symbols-public/v1/",
        'hang_profile_filename': 'hang_profile_128_16000',
        'print_debug_info': False,
        'hang_lower_bound': 128,
        'hang_upper_bound': 16000,
        'stack_acceptance_threshold': 0.001,
        'hang_outlier_threshold': 512,
        'uuid': uuid.uuid4().hex,
    }

    if config is not None:
        final_config.update(config)

    iterations = (final_config['days_to_aggregate'] + final_config['date_clumping'] - 1) / final_config['date_clumping']
    profile_processor = ProfileProcessor(final_config)

    job_start = time.time()
    # We were OOMing trying to allocate a contiguous array for all of this. Pass it in
    # bit by bit to the profile processor and hope it can handle it.
    for x in xrange(0, iterations):
        day_start = time.time()
        rel_end = x * final_config['date_clumping']
        rel_end = -rel_end - 1
        rel_start = min((x + 1) * final_config['date_clumping'] - 1, final_config['days_to_aggregate'])
        rel_start = -rel_start - 1
        data = time_code("Getting data",
            lambda: get_data(sc, final_config, rel_start, rel_end))
        if data is None:
            continue
        transformed = transform_pings(sc, data, final_config)
        time_code("Passing stacks to processor", lambda: profile_processor.ingest(transformed))
        # Run a collection to ensure that any references to any RDDs are cleaned up,
        # allowing the JVM to clean them up on its end.
        gc.collect()
        day_end = time.time()
        day_delta = day_end - day_start
        print "Iteration took {}s".format(int(round(day_delta)))
        job_elapsed = time.time() - job_start
        percent_done = float(x + 1) / float(iterations)
        projected = job_elapsed / percent_done
        remaining = projected - job_elapsed
        print "Job should finish in {}".format(timedelta(seconds=remaining))

    profile = profile_processor.process_into_profile()
    write_file(final_config['hang_profile_filename'], profile, final_config)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Electrum - lightweight Bitcoin client
# Copyright (C) 2012 thomasv@ecdsa.org
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import util
import bitcoin
from bitcoin import *

MAX_TARGET = 0x00000000FFFF0000000000000000000000000000000000000000000000000000

def serialize_header(res):
    s = int_to_hex(res.get('version'), 4) \
        + rev_hex(res.get('prev_block_hash')) \
        + rev_hex(res.get('merkle_root')) \
        + int_to_hex(int(res.get('timestamp')), 4) \
        + int_to_hex(int(res.get('bits')), 4) \
        + int_to_hex(int(res.get('nonce')), 4)
    return s

def deserialize_header(s, height):
    hex_to_int = lambda s: int('0x' + s[::-1].encode('hex'), 16)
    h = {}
    h['version'] = hex_to_int(s[0:4])
    h['prev_block_hash'] = hash_encode(s[4:36])
    h['merkle_root'] = hash_encode(s[36:68])
    h['timestamp'] = hex_to_int(s[68:72])
    h['bits'] = hex_to_int(s[72:76])
    h['nonce'] = hex_to_int(s[76:80])
    h['block_height'] = height
    return h

def hash_header(header):
    if header is None:
        return '0' * 64
    if header.get('prev_block_hash') is None:
        header['prev_block_hash'] = '00'*32
    return hash_encode(Hash(serialize_header(header).decode('hex')))

blockchains = {}

def read_blockchains(config):
    blockchains[0] = Blockchain(config, 'blockchain_headers')
    l = filter(lambda x: x.startswith('fork_'), os.listdir(config.path))
    l = sorted(l, key = lambda x: int(x.split('_')[1]))
    for x in l:
        b = Blockchain(config, x)
        blockchains[b.checkpoint] = b
    return blockchains

def check_header(header):
    if type(header) is not dict:
        return False
    for b in blockchains.values():
        if b.check_header(header):
            return b
    return False

def can_connect(header):
    for b in blockchains.values():
        if b.can_connect(header):
            return b
    return False

class Blockchain(util.PrintError):

    '''Manages blockchain headers and their verification'''

    def __init__(self, config, filename):
        self.config = config
        self.filename = filename
        self.catch_up = None # interface catching up
        self.is_saved = True
        self.headers = []
        if filename == 'blockchain_headers':
            self.parent = None
            self.checkpoint = 0
        elif filename.startswith('fork_'):
            self.parent = blockchains[int(filename.split('_')[1])]
            self.checkpoint = int(filename.split('_')[2])
        else:
            raise BaseException('')

    def get_max_child(self):
        children = filter(lambda y: y.parent==self, blockchains.values())
        return max([x.checkpoint for x in children]) if children else None

    def get_branch_size(self):
        mc = self.get_max_child()
        checkpoint = mc if mc is not None else self.checkpoint
        return self.height() - checkpoint

    def check_header(self, header):
        header_hash = hash_header(header)
        height = header.get('block_height')
        return header_hash == self.get_hash(height)

    def fork(parent, checkpoint):
        filename = 'fork_%d_%d'%(parent.checkpoint, checkpoint)
        self = Blockchain(parent.config, filename)
        self.is_saved = False
        return self

    def height(self):
        return self.checkpoint + self.size() - 1

    def size(self):
        if self.is_saved:
            p = self.path()
            return os.path.getsize(p)/80 if os.path.exists(p) else 0
        else:
            return len(self.headers)

    def verify_header(self, header, prev_header, bits, target):
        prev_hash = hash_header(prev_header)
        _hash = hash_header(header)
        if prev_hash != header.get('prev_block_hash'):
            raise BaseException("prev hash mismatch: %s vs %s" % (prev_hash, header.get('prev_block_hash')))
        if bitcoin.TESTNET:
            return
        if bits != header.get('bits'):
            raise BaseException("bits mismatch: %s vs %s" % (bits, header.get('bits')))
        if int('0x' + _hash, 16) > target:
            raise BaseException("insufficient proof of work: %s vs target %s" % (int('0x' + _hash, 16), target))

    def verify_chain(self, chain):
        first_header = chain[0]
        prev_header = self.read_header(first_header.get('block_height') - 1)
        for header in chain:
            height = header.get('block_height')
            bits, target = self.get_target(height / 2016, chain)
            self.verify_header(header, prev_header, bits, target)
            prev_header = header

    def verify_chunk(self, index, data):
        num = len(data) / 80
        prev_header = None
        if index != 0:
            prev_header = self.read_header(index*2016 - 1)
        bits, target = self.get_target(index)
        for i in range(num):
            raw_header = data[i*80:(i+1) * 80]
            header = deserialize_header(raw_header, index*2016 + i)
            self.verify_header(header, prev_header, bits, target)
            prev_header = header

    def path(self):
        d = util.get_headers_dir(self.config)
        return os.path.join(d, self.filename)

    def save_chunk(self, index, chunk):
        if not self.is_saved:
            self.save()
        filename = self.path()
        d = (index * 2016 - self.checkpoint) * 80
        if d < 0:
            chunk = chunk[-d:]
            d = 0
        with open(filename, 'rb+') as f:
            f.seek(d)
            f.write(chunk)

    def save(self):
        # recursively save parents if they have not been saved
        if self.parent and not self.parent.is_saved:
            self.parent.save()
        open(self.path(), 'w+').close()
        for h in self.headers:
            self.write_header(h)
        self.headers = []
        self.is_saved = True
        self.print_error("saved", self.filename)

    def swap_with_parent(self):
        self.print_error("swap")
        parent = self.parent
        checkpoint = self.checkpoint
        # copy headers
        parent.headers = [parent.read_header(h) for h in range(checkpoint, checkpoint + parent.get_branch_size())]
        # truncate parent file
        with open(parent.path(), 'rb+') as f:
            f.seek(checkpoint*80)
            f.truncate()
        parent.is_saved = False
        # swap chains
        fn = self.filename; self.filename = parent.filename; parent.filename = fn
        self.parent = parent.parent; parent.parent = parent
        self.checkpoint = parent.checkpoint; parent.checkpoint = checkpoint
        # write my headers
        for h in self.headers:
            self.write_header(h)
        self.headers = []
        self.is_saved = True

    def save_header(self, header):
        N = 10
        height = header.get('block_height')
        if not self.is_saved:
            assert height == self.checkpoint + len(self.headers)
            self.headers.append(header)
            if len(self.headers) > N:
                if self.parent.get_branch_size() <= N:
                    self.swap_with_parent()
                else:
                    self.save()
            return
        self.write_header(header)

    def write_header(self, header):
        filename = self.path()
        delta = header.get('block_height') - self.checkpoint
        data = serialize_header(header).decode('hex')
        assert delta * 80 == os.path.getsize(filename)
        assert len(data) == 80
        with open(filename, 'rb+') as f:
            f.seek(delta * 80)
            f.write(data)

    def read_header(self, height):
        if height < self.checkpoint:
            return self.parent.read_header(height)
        delta = height - self.checkpoint
        if not self.is_saved:
            if delta >= len(self.headers):
                return None
            header = self.headers[delta]
            assert header.get('block_height') == height
            return header
        name = self.path()
        if os.path.exists(name):
            f = open(name, 'rb')
            f.seek(delta * 80)
            h = f.read(80)
            f.close()
            if len(h) == 80:
                h = deserialize_header(h, height)
                return h

    def get_hash(self, height):
        return bitcoin.GENESIS if height == 0 else hash_header(self.read_header(height))

    def BIP9(self, height, flag):
        v = self.read_header(height)['version']
        return ((v & 0xE0000000) == 0x20000000) and ((v & flag) == flag)

    def segwit_support(self, N=144):
        h = self.local_height
        return sum([self.BIP9(h-i, 2) for i in range(N)])*10000/N/100.

    def truncate_headers(self, height):
        self.print_error('Truncating headers file at height %d'%height)
        name = self.path()
        f = open(name, 'rb+')
        f.seek(height * 80)
        f.truncate()
        f.close()

    def get_target(self, index, chain=None):
        if bitcoin.TESTNET:
            return 0, 0
        if index == 0:
            return 0x1d00ffff, MAX_TARGET
        first = self.read_header((index-1) * 2016)
        last = self.read_header(index*2016 - 1)
        if last is None:
            for h in chain:
                if h.get('block_height') == index*2016 - 1:
                    last = h
        assert last is not None
        # bits to target
        bits = last.get('bits')
        bitsN = (bits >> 24) & 0xff
        if not (bitsN >= 0x03 and bitsN <= 0x1d):
            raise BaseException("First part of bits should be in [0x03, 0x1d]")
        bitsBase = bits & 0xffffff
        if not (bitsBase >= 0x8000 and bitsBase <= 0x7fffff):
            raise BaseException("Second part of bits should be in [0x8000, 0x7fffff]")
        target = bitsBase << (8 * (bitsN-3))
        # new target
        nActualTimespan = last.get('timestamp') - first.get('timestamp')
        nTargetTimespan = 14 * 24 * 60 * 60
        nActualTimespan = max(nActualTimespan, nTargetTimespan / 4)
        nActualTimespan = min(nActualTimespan, nTargetTimespan * 4)
        new_target = min(MAX_TARGET, (target*nActualTimespan) / nTargetTimespan)
        # convert new target to bits
        c = ("%064x" % new_target)[2:]
        while c[:2] == '00' and len(c) > 6:
            c = c[2:]
        bitsN, bitsBase = len(c) / 2, int('0x' + c[:6], 16)
        if bitsBase >= 0x800000:
            bitsN += 1
            bitsBase >>= 8
        new_bits = bitsN << 24 | bitsBase
        return new_bits, bitsBase << (8 * (bitsN-3))

    def can_connect(self, header):
        previous_height = header['block_height'] - 1
        previous_header = self.read_header(previous_height)
        if not previous_header:
            return False
        prev_hash = hash_header(previous_header)
        if prev_hash != header.get('prev_block_hash'):
            return False
        height = header.get('block_height')
        bits, target = self.get_target(height / 2016)
        try:
            self.verify_header(header, previous_header, bits, target)
        except:
            return False
        return True

    def connect_chunk(self, idx, hexdata):
        try:
            data = hexdata.decode('hex')
            self.verify_chunk(idx, data)
            #self.print_error("validated chunk %d" % idx)
            self.save_chunk(idx, data)
            return True
        except BaseException as e:
            self.print_error('verify_chunk failed', str(e))
            return False
<commit_msg>test height<commit_after>#!/usr/bin/env python
#
# Electrum - lightweight Bitcoin client
# Copyright (C) 2012 thomasv@ecdsa.org
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import util
import bitcoin
from bitcoin import *

MAX_TARGET = 0x00000000FFFF0000000000000000000000000000000000000000000000000000

def serialize_header(res):
    s = int_to_hex(res.get('version'), 4) \
        + rev_hex(res.get('prev_block_hash')) \
        + rev_hex(res.get('merkle_root')) \
        + int_to_hex(int(res.get('timestamp')), 4) \
        + int_to_hex(int(res.get('bits')), 4) \
        + int_to_hex(int(res.get('nonce')), 4)
    return s

def deserialize_header(s, height):
    hex_to_int = lambda s: int('0x' + s[::-1].encode('hex'), 16)
    h = {}
    h['version'] = hex_to_int(s[0:4])
    h['prev_block_hash'] = hash_encode(s[4:36])
    h['merkle_root'] = hash_encode(s[36:68])
    h['timestamp'] = hex_to_int(s[68:72])
    h['bits'] = hex_to_int(s[72:76])
    h['nonce'] = hex_to_int(s[76:80])
    h['block_height'] = height
    return h

def hash_header(header):
    if header is None:
        return '0' * 64
    if header.get('prev_block_hash') is None:
        header['prev_block_hash'] = '00'*32
    return hash_encode(Hash(serialize_header(header).decode('hex')))

blockchains = {}

def read_blockchains(config):
    blockchains[0] = Blockchain(config, 'blockchain_headers')
    l = filter(lambda x: x.startswith('fork_'), os.listdir(config.path))
    l = sorted(l, key = lambda x: int(x.split('_')[1]))
    for x in l:
        b = Blockchain(config, x)
        blockchains[b.checkpoint] = b
    return blockchains

def check_header(header):
    if type(header) is not dict:
        return False
    for b in blockchains.values():
        if b.check_header(header):
            return b
    return False

def can_connect(header):
    for b in blockchains.values():
        if b.can_connect(header):
            return b
    return False

class Blockchain(util.PrintError):

    '''Manages blockchain headers and their verification'''

    def __init__(self, config, filename):
        self.config = config
        self.filename = filename
        self.catch_up = None # interface catching up
        self.is_saved = True
        self.headers = []
        if filename == 'blockchain_headers':
            self.parent = None
            self.checkpoint = 0
        elif filename.startswith('fork_'):
            self.parent = blockchains[int(filename.split('_')[1])]
            self.checkpoint = int(filename.split('_')[2])
        else:
            raise BaseException('')

    def get_max_child(self):
        children = filter(lambda y: y.parent==self, blockchains.values())
        return max([x.checkpoint for x in children]) if children else None

    def get_branch_size(self):
        mc = self.get_max_child()
        checkpoint = mc if mc is not None else self.checkpoint
        return self.height() - checkpoint

    def check_header(self, header):
        header_hash = hash_header(header)
        height = header.get('block_height')
        return header_hash == self.get_hash(height)

    def fork(parent, checkpoint):
        filename = 'fork_%d_%d'%(parent.checkpoint, checkpoint)
        self = Blockchain(parent.config, filename)
        self.is_saved = False
        return self

    def height(self):
        return self.checkpoint + self.size() - 1

    def size(self):
        if self.is_saved:
            p = self.path()
            return os.path.getsize(p)/80 if os.path.exists(p) else 0
        else:
            return len(self.headers)

    def verify_header(self, header, prev_header, bits, target):
        prev_hash = hash_header(prev_header)
        _hash = hash_header(header)
        if prev_hash != header.get('prev_block_hash'):
            raise BaseException("prev hash mismatch: %s vs %s" % (prev_hash, header.get('prev_block_hash')))
        if bitcoin.TESTNET:
            return
        if bits != header.get('bits'):
            raise BaseException("bits mismatch: %s vs %s" % (bits, header.get('bits')))
        if int('0x' + _hash, 16) > target:
            raise BaseException("insufficient proof of work: %s vs target %s" % (int('0x' + _hash, 16), target))

    def verify_chain(self, chain):
        first_header = chain[0]
        prev_header = self.read_header(first_header.get('block_height') - 1)
        for header in chain:
            height = header.get('block_height')
            bits, target = self.get_target(height / 2016, chain)
            self.verify_header(header, prev_header, bits, target)
            prev_header = header

    def verify_chunk(self, index, data):
        num = len(data) / 80
        prev_header = None
        if index != 0:
            prev_header = self.read_header(index*2016 - 1)
        bits, target = self.get_target(index)
        for i in range(num):
            raw_header = data[i*80:(i+1) * 80]
            header = deserialize_header(raw_header, index*2016 + i)
            self.verify_header(header, prev_header, bits, target)
            prev_header = header

    def path(self):
        d = util.get_headers_dir(self.config)
        return os.path.join(d, self.filename)

    def save_chunk(self, index, chunk):
        if not self.is_saved:
            self.save()
        filename = self.path()
        d = (index * 2016 - self.checkpoint) * 80
        if d < 0:
            chunk = chunk[-d:]
            d = 0
        with open(filename, 'rb+') as f:
            f.seek(d)
            f.write(chunk)

    def save(self):
        # recursively save parents if they have not been saved
        if self.parent and not self.parent.is_saved:
            self.parent.save()
        open(self.path(), 'w+').close()
        for h in self.headers:
            self.write_header(h)
        self.headers = []
        self.is_saved = True
        self.print_error("saved", self.filename)

    def swap_with_parent(self):
        self.print_error("swap")
        parent = self.parent
        checkpoint = self.checkpoint
        # copy headers
        parent.headers = [parent.read_header(h) for h in range(checkpoint, checkpoint + parent.get_branch_size())]
        # truncate parent file
        with open(parent.path(), 'rb+') as f:
            f.seek(checkpoint*80)
            f.truncate()
        parent.is_saved = False
        # swap chains
        fn = self.filename; self.filename = parent.filename; parent.filename = fn
        self.parent = parent.parent; parent.parent = parent
        self.checkpoint = parent.checkpoint; parent.checkpoint = checkpoint
        # write my headers
        for h in self.headers:
            self.write_header(h)
        self.headers = []
        self.is_saved = True

    def save_header(self, header):
        N = 10
        height = header.get('block_height')
        if not self.is_saved:
            assert height == self.checkpoint + len(self.headers)
            self.headers.append(header)
            if len(self.headers) > N:
                if self.parent.get_branch_size() <= N:
                    self.swap_with_parent()
                else:
                    self.save()
            return
        self.write_header(header)

    def write_header(self, header):
        filename = self.path()
        delta = header.get('block_height') - self.checkpoint
        data = serialize_header(header).decode('hex')
        assert delta * 80 == os.path.getsize(filename)
        assert len(data) == 80
        with open(filename, 'rb+') as f:
            f.seek(delta * 80)
            f.write(data)

    def read_header(self, height):
        if height < self.checkpoint:
            return self.parent.read_header(height)
        delta = height - self.checkpoint
        if not self.is_saved:
            if delta >= len(self.headers):
                return None
            header = self.headers[delta]
            assert header.get('block_height') == height
            return header
        name = self.path()
        if os.path.exists(name):
            f = open(name, 'rb')
            f.seek(delta * 80)
            h = f.read(80)
            f.close()
            if len(h) == 80:
                h = deserialize_header(h, height)
                return h

    def get_hash(self, height):
        return bitcoin.GENESIS if height == 0 else hash_header(self.read_header(height))

    def BIP9(self, height, flag):
        v = self.read_header(height)['version']
        return ((v & 0xE0000000) == 0x20000000) and ((v & flag) == flag)

    def segwit_support(self, N=144):
        h = self.local_height
        return sum([self.BIP9(h-i, 2) for i in range(N)])*10000/N/100.

    def truncate_headers(self, height):
        self.print_error('Truncating headers file at height %d'%height)
        name = self.path()
        f = open(name, 'rb+')
        f.seek(height * 80)
        f.truncate()
        f.close()

    def get_target(self, index, chain=None):
        if bitcoin.TESTNET:
            return 0, 0
        if index == 0:
            return 0x1d00ffff, MAX_TARGET
        first = self.read_header((index-1) * 2016)
        last = self.read_header(index*2016 - 1)
        if last is None:
            for h in chain:
                if h.get('block_height') == index*2016 - 1:
                    last = h
        assert last is not None
        # bits to target
        bits = last.get('bits')
        bitsN = (bits >> 24) & 0xff
        if not (bitsN >= 0x03 and bitsN <= 0x1d):
            raise BaseException("First part of bits should be in [0x03, 0x1d]")
        bitsBase = bits & 0xffffff
        if not (bitsBase >= 0x8000 and bitsBase <= 0x7fffff):
            raise BaseException("Second part of bits should be in [0x8000, 0x7fffff]")
        target = bitsBase << (8 * (bitsN-3))
        # new target
        nActualTimespan = last.get('timestamp') - first.get('timestamp')
        nTargetTimespan = 14 * 24 * 60 * 60
        nActualTimespan = max(nActualTimespan, nTargetTimespan / 4)
        nActualTimespan = min(nActualTimespan, nTargetTimespan * 4)
        new_target = min(MAX_TARGET, (target*nActualTimespan) / nTargetTimespan)
        # convert new target to bits
        c = ("%064x" % new_target)[2:]
        while c[:2] == '00' and len(c) > 6:
            c = c[2:]
        bitsN, bitsBase = len(c) / 2, int('0x' + c[:6], 16)
        if bitsBase >= 0x800000:
            bitsN += 1
            bitsBase >>= 8
        new_bits = bitsN << 24 | bitsBase
        return new_bits, bitsBase << (8 * (bitsN-3))

    def can_connect(self, header):
        height = header['block_height']
        if self.height() != height - 1:
            return False
        previous_header = self.read_header(height -1)
        if not previous_header:
            return False
        prev_hash = hash_header(previous_header)
        if prev_hash != header.get('prev_block_hash'):
            return False
        bits, target = self.get_target(height / 2016)
        try:
            self.verify_header(header, previous_header, bits, target)
        except:
            return False
        return True

    def connect_chunk(self, idx, hexdata):
        try:
            data = hexdata.decode('hex')
            self.verify_chunk(idx, data)
            #self.print_error("validated chunk %d" % idx)
            self.save_chunk(idx, data)
            return True
        except BaseException as e:
            self.print_error('verify_chunk failed', str(e))
            return False
<|endoftext|>