import json
import org.parosproxy.paros.network.HttpRequestHeader as HttpRequestHeader
import org.parosproxy.paros.network.HttpHeader as HttpHeader
import org.zaproxy.zap.extension.script.ScriptVars as GlobalVariables
import org.parosproxy.paros.network.HttpMessage as HttpMessage
import org.parosproxy.paros.network.HtmlParameter as HtmlParameter
from org.apache.commons.httpclient import URI
 
import java.util.concurrent.TimeUnit as TimeUnit
import org.parosproxy.paros.control.Control as Control
import org.zaproxy.zap.extension.selenium.ExtensionSelenium as ExtensionSelenium
import org.openqa.selenium.By as By
import time
import threading
import org.openqa.selenium.chrome.ChromeDriver as ChromeDriver
import org.openqa.selenium.firefox.FirefoxDriver as FirefoxDriver;
import org.openqa.selenium.firefox.FirefoxOptions as FirefoxOptions;
import org.openqa.selenium.JavascriptExecutor as JavascriptExecutor;
import org.openqa.selenium.WebElement as WebElement;
from synchronize import make_synchronized
 
# base url of the juice shop web app
# TODO: change this to the url of your Juice Shop
BASE_URL = 'http://localhost:3000';
 
# TODO: change this to your own google username and password
YOUR_GOOGLE_EMAIL = '<YOUR GOOGLE EMAIL>'
YOUR_GOOGLE_PASSWORD = '<YOUR GOOGLE PASSWORD>'
 
# function called for every outgoing request to juice shop (part of httpsender)
def sendingRequest(msg, initiator, helper):
    print('sendingRequest called for url=' + msg.getRequestHeader().getURI().toString())
 
    #checking if we already have an access token
    accessToken = GlobalVariables.getGlobalVar("accessToken");    
    if accessToken is not None:
        print "we have access token, checking if token is valid";
        if isTokenValid(accessToken, helper) == True:
            print "accessToken in valid";
            setAccessTokenInHttpMessage(accessToken, msg);
            return;
 
    print "token is invalid or there is no token, need to relogin"
    accessToken = refreshAccessToken(helper);
    print "accessToken:"+accessToken;
    setAccessTokenInHttpMessage(accessToken, msg);
    return;
 
#we have to make this function synchronized as we do not want to have duplicate concurrent attempts to login
@make_synchronized
def refreshAccessToken(helper):
    print "refreshing access token and checking if it has already been refreshed"
    accessToken = GlobalVariables.getGlobalVar("accessToken");    
    if accessToken is not None and isTokenValid(accessToken, helper) == True:
        print "access token already refreshed, no need to relogin"
        return accessToken;
 
    GlobalVariables.setGlobalVar("accessToken", None);
    accessToken = doLogin(helper);
    GlobalVariables.setGlobalVar("accessToken", str(accessToken)); 
    print "access token refreshed"
    return accessToken;
 
def isTokenValid(accessToken, helper):
    whoamiUrl = BASE_URL + "/rest/user/whoami";    
    returnedMsg = callGet(helper, whoamiUrl, accessToken);
    userJsonObject = json.loads(str(returnedMsg.getResponseBody()));
    userId = userJsonObject.get("user").get("id");
    print "userId returned:"+str(userId);
    if userId is None:
        return False;
 
    return True;
 
#this is where we do the login using google sso 
def doLogin(helper):
    firefoxOptions = FirefoxOptions()
    firefoxOptions.addArguments("--window-size=1920,1080");
    firefoxOptions.addArguments("--disable-gpu");
    firefoxOptions.addArguments("--disable-extensions");
    firefoxOptions.addArguments("--proxy-server='direct://'");
    firefoxOptions.addArguments("--proxy-bypass-list=*");
    firefoxOptions.addArguments("--start-maximized");
    firefoxOptions.addArguments("--headless");
    webDriver = FirefoxDriver(firefoxOptions);
    webDriver.get(BASE_URL);
 
    #clicking on the Dismiss button to dismiss the popup dialog box. The css selector could be improved so that it is not hardcoded to use div:nth-child(3)
    # but will do for now
    ele = webDriver.findElement(By.cssSelector("#mat-dialog-0 > app-welcome-banner > div > div:nth-child(3) > button.mat-focus-indicator.close-dialog.mat-raised-button.mat-button-base.mat-primary.ng-star-inserted"));
    click(ele, webDriver);
 
    #clicking on the Account link    
    webDriver.manage().timeouts().implicitlyWait(3, TimeUnit.SECONDS);
    ele = webDriver.findElement(By.cssSelector("#navbarAccount > span.mat-button-wrapper > span"));
    click(ele, webDriver);
 
    #clicking on the Login link    
    ele = webDriver.findElement(By.cssSelector("#navbarLoginButton > span"));
    click(ele, webDriver);
 
    #clicking on the login using google sso link 
    ele = webDriver.findElement(By.cssSelector("#loginButtonGoogle > span.mat-button-wrapper"));
    click(ele, webDriver);
 
    #wait for google page to load
    webDriver.manage().timeouts().implicitlyWait(2, TimeUnit.SECONDS);
 
    #keying in your google username
    webDriver.findElement(By.id("identifierId")).sendKeys(YOUR_GOOGLE_EMAIL);
 
    #clicking on next
    ele = webDriver.findElement(By.cssSelector(".VfPpkd-LgbsSe-OWXEXe-k8QpJ > div:nth-child(3)"));
    click(ele, webDriver);
 
    #wait for page to load
    time.sleep(2);#in seconds
 
    #key in password
    password = webDriver.findElement(By.name("password"));
    password.sendKeys(YOUR_GOOGLE_PASSWORD);
 
    #submit the login credentials
    ele = webDriver.findElement(By.cssSelector(".VfPpkd-LgbsSe-OWXEXe-k8QpJ > span:nth-child(2)"));
    click(ele, webDriver);
    time.sleep(4);#in seconds
    accessToken = getAccessToken(webDriver);
    return accessToken;
 
def click(ele, webDriver):
    webDriver.executeScript("arguments[0].click()", ele);
 
# what we are doing here is to wait for the browser to arrive back at the Juice Shop page, at which point the cookie holding the token will be available
def getAccessToken(webDriver):
    accessTokenCookie = None;
    for number in range(5):
        accessTokenCookie = webDriver.manage().getCookieNamed("token");
        if accessTokenCookie is None:
            print "accessTokenCookie is None, sleeping for a while more";
            time.sleep(5)
        else:
            break
 
    return accessTokenCookie.getValue();
 
 
# function called for every incoming server response from juice shop (part of httpsender)
def responseReceived(msg, initiator, helper):
    pass
 
# function to make generic GET request to Juice Shop
def callGet(helper, requestUrl, accessToken):
    requestUri = URI(requestUrl, False);
    print "-----start of callGet-------";
    print "requestUrl:"+requestUrl;
    msg = HttpMessage();
    requestHeader = HttpRequestHeader(HttpRequestHeader.GET, requestUri, HttpHeader.HTTP10);
    msg.setRequestHeader(requestHeader);
    setAccessTokenInHttpMessage(accessToken, msg);
    print "Sending GET request: " + str(requestHeader);
    helper.getHttpSender().sendAndReceive(msg)
    print "Received response status code for authentication request: " + str(msg.getResponseHeader());
    print "------------------------------------";
    return msg;
 
# function to set the token in Authorization header and in cookie in request
def setAccessTokenInHttpMessage(accessToken, msg):
    print "setting token in request"
    msg.getRequestHeader().setHeader("Authorization", "Bearer " + accessToken);
    cookie = HtmlParameter(HtmlParameter.Type.cookie, "token", accessToken);
    cookies = msg.getRequestHeader().getCookieParams();
    cookies.add(cookie);
    msg.getRequestHeader().setCookieParams(cookies);