
import keyring
import requests
import json
import logging
import html

#import urllib3
#urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

keyring.get_keyring()

### Define global variables
MAX_CX_RESULTS = 10
CXSAST_RESULTS_XML = "WebGoat.Net.xml"
SERVICE_NAME = "TFS2017"
API_VERSION = "4.1"
USER_NAME = "tyler"
TFS2017_TOKEN = keyring.get_password(SERVICE_NAME, USER_NAME)
TFS2017_ORGANIZATION = ""
TFS2017_URL = "https://tylera-laptop:444/tfs/DefaultCollection/"
PROJECT_NAME = "WebGoat .NET"
SIMILARITY_ID_FIELD ='Checkmarx.SimilarityID'
NODE_ID_FIELD ='Checkmarx.NodeID'
IGNORE_FP = True
HEADERS = {"Authorization": "Basic " + TFS2017_TOKEN,'Content-type': 'application/json'}


### Define Issue class for assistance with creating/adding issues to an array. Pretty print method for debugging ###
class Issue:
	def __init__(self, filename='', deeplink='', description='', severity='', similarityID='',pathNodeID='',wid=0, vulnerability='', assignedToUser=''):
		self.filename = filename
		self.deeplink = deeplink
		self.description = description
		self.severity = severity
		self.similarityID = similarityID
		self.pathNodeID = pathNodeID
		self.wid = wid
		self.vulnerability = vulnerability
		self.assignedToUser = assignedToUser

	def __eq__(self, other):
		return (self.similarityID + self.pathNodeID == other.similarityID + other.pathNodeID)
	def __str__(self):
		return ("WID: " + str(self.wid) + 
		"\nPathNodeID : " + str(self.pathNodeID) + 
		"\nSimilarityId # : " + str(self.similarityID) + 
		"\nFileName: " + self.filename + 
		"\nDeepLink: " + self.deeplink +
		"\nDescription: " + self.description + 
		"\nSeverity: " + self.severity +
		"\nVulnerability: " + self.vulnerability +
		"\nAssignedTo: " + self.assignedToUser)

### HELPER METHODS ###
### Will parse out the short file name from full path.
def prepFileName(filename):
	sourceList = filename.split('/')
	return sourceList[len(sourceList)-1]
### Parse XML from post-action file generated by Checkmarx SAST scan. Look for similarity IDs with Status = 'New' 
### XML Structure - CxXMLResults -> Query -> Result -> Path -> PathNode
def parse_xml(xmlFilePath):
	listOfIssues = list()
	import xml.etree.ElementTree as ET
	tree = ET.parse(xmlFilePath)
	root = tree.getroot()
	print("ProjectName: " + root.attrib["ProjectName"])
	print("InitiatorName: " + root.attrib["InitiatorName"])
	#Loop through Query tags
	queryResultCount = 0
	
	for query in root:
		vulnerability = query.attrib["name"]
		for result in query:
			if (IGNORE_FP == True and result.attrib["FalsePositive"] != "True") or (IGNORE_FP == False):
				if queryResultCount < MAX_CX_RESULTS:
					#Idenfity "New" issues
					if(result.attrib["Severity"] == "High"):
						for path in result:
							description = "Source " + path[0][0].text + " Line " + path[0][1].text + " Object " + path[0][4].text + " "
							description += "Destination " + path[len(path)-1][0].text + " Line " + path[len(path)-1][1].text + " Object " + path[len(path)-1][4].text + " "
							description += r'<a href= "' + result.attrib["DeepLink"] + r'">' + result.attrib["DeepLink"] + r'</a>'
							issue = Issue(result.attrib["FileName"], result.attrib["DeepLink"],description, result.attrib["Severity"], path.attrib["SimilarityId"], str(result.attrib["NodeId"]), 0, vulnerability, result.attrib["AssignToUser"])
							listOfIssues.append(issue)
						queryResultCount += 1

	if len(listOfIssues) > 0:
		return listOfIssues
	else:
		### There are no HV in the list
		return None

### Call TFS2017 API to utilize custom query to get all workItems from a specific project
def get_tfs2017_workitems():
	headers = {
		"Authorization": "Basic " + TFS2017_TOKEN
	}
	query = r"SELECT System.ID, System.Title FROM workitems WHERE System.State <> 'RESOLVED' AND " + SIMILARITY_ID_FIELD + r" <> '' ORDER BY System.ID ASC"
	json = 	{"query": query}

	url = TFS2017_URL + "_apis/wit/wiql?api-version=" + API_VERSION
	tfs2017_response = requests.post(url, json=json, headers=HEADERS, verify=False)

	if tfs2017_response.status_code == 200:
		response = tfs2017_response.json()
		if response["workItems"]:
			### String to hold all of the Work Item IDS for the next API call
			wid = ""
			### Loop through workItems to pull the ids
			for work_item in response["workItems"]:
				### Make sure the id is an int greater than 0
				if isinstance(work_item["id"], int) and work_item["id"] > 0:
					### Concatentate the id to the wid string
					### if an id already exists separate by a comma
					if len(wid) > 0:
						wid = wid + "," + str(work_item["id"])
					else:
						wid = str(work_item["id"])    
			### If wid is empty return None
			if len(wid) < 1:
				return None

			### Call TFS2017 API to get work item details based on the ids loaded into wid
			### https://docs.microsoft.com/en-us/rest/api/tfs2017/wit/work%20items/list?view=tfs2017-rest-4.1
			url = TFS2017_URL + '_apis/wit/workitems?ids=' + wid + '&api-version=' + API_VERSION
			newtfs2017 = requests.get(url, headers=HEADERS, verify=False)    

			if newtfs2017.status_code == 200: 
				work_item_response = newtfs2017.json()
				workItemList = list()
				witems = work_item_response['value']
				for witem in witems:
					if(SIMILARITY_ID_FIELD in witem['fields']):
						workItem = Issue(similarityID= witem['fields'][SIMILARITY_ID_FIELD],pathNodeID= witem['fields'][NODE_ID_FIELD],wid = witem['id'])
						workItemList.append(workItem)
					# print(workItemList)
				return workItemList
			else:
				return None
	else:
		return None

def find_new_issues(xmlList,witemList):
	newIssues = list()
	for xml in xmlList:
		notFound = 0
		if witemList is not None:
			for witem in witemList:
				if witem == xml:
					notFound = 1
		if notFound == 0:
			newIssues.append(xml)
	return newIssues

def find_resolved_issues(xmlList,witemList):
	resolvedIssues = list()
	if witemList is not None:
		for witem in witemList:
			isFound = 0
			for xml in xmlList:
				if witem == xml:
					isFound = 1
			if isFound == 0:
				resolvedIssues.append(witem)
	return resolvedIssues

def addNewIssue(issue):
	filename = prepFileName(issue.filename)
	headers = {
		"Authorization": "Basic " + TFS2017_TOKEN,
		'Content-type': 'application/json-patch+json'
	}
	issueData = [
		{
			"op": "add",
			"path": "/fields/System.Title",
			"value": issue.vulnerability + " - " + issue.similarityID + " - " + filename
		},
		{
			"op": "add",
			"path": "/fields/Microsoft.VSTS.Common.Priority",
			"value": 2
		},
		{
			"op": "add",
			"path": "/fields/Microsoft.VSTS.Common.Severity",
			"value": "2 - High"
		},
		{
			"op": "add",
			"path": "/fields/" + SIMILARITY_ID_FIELD,
			"value": issue.similarityID
		},
		{
			"op": "add",
			"path": "/fields/" + NODE_ID_FIELD,
			"value": issue.pathNodeID
		},
		{
			"op": "add",
			"path": "/fields/Microsoft.VSTS.TCM.ReproSteps",
			"value": issue.description
		},
		{
			"op": "add",
			"path": "/fields/System.WorkItemType",
			"value": "Bug"
		}
	]

	postNewIssue = requests.patch(TFS2017_URL + PROJECT_NAME + r'/_apis/wit/workitems/$bug?api-version=' + API_VERSION, json=issueData, headers=headers, verify=False)   
	if postNewIssue.status_code == 200: 
		print("New Bugs Posted")
	else:
		print(postNewIssue.status_code)

def resolve_issues(witemList):
	headers = {
		"Authorization": "Basic " + TFS2017_TOKEN,
		'Content-type': 'application/json-patch+json'
	}
	issueData = [
		{
			"op": "add",
			"path": "/fields/System.State",
			"value": "Resolved"
		}
	]
		
	if len(witemList) > 0:
		for witem in witemList:
			# print(witem)
			resolveIssue = requests.patch(TFS2017_URL + PROJECT_NAME + '/_apis/wit/workitems/'+ str(witem.wid) + '?api-version=' + API_VERSION, json=issueData, headers=headers, verify = False)   
			if resolveIssue.status_code == 200: 
				print("Bug Resolution Successful!")
			else:
				print(resolveIssue.status_code)
	else:
		print("There are no resolved issues")

### Get Work Items List
witems = get_tfs2017_workitems()

### Xml List
xml = parse_xml(CXSAST_RESULTS_XML)

### Compare XML and Work Items List to find new issues
newIssues = find_new_issues(xml,witems)

### Loop through new issues and add them into tfs2017
for issue in newIssues:
	addNewIssue(issue)

### Compare CXSAST XML and Work Items List to find issues that are no longer in the CX XML Results, but are still in Work Items
resolvedIssuesList = find_resolved_issues(xml,witems)

### Mark issues that no longer exist as resolved
resolve_issues(resolvedIssuesList)