# Simple HTTP example
#
# A simple example using the HTTP plugin that shows the retrieval of a
# single page via HTTP. The resulting page is written to a file.
#
# More complex HTTP scripts are best created with the TCPProxy.

from net.grinder.script.Grinder import grinder
from net.grinder.script import Test
from net.grinder.plugin.http import HTTPRequest
from net.grinder.plugin.http import HTTPPluginControl
from HTTPClient import CookieModule
from HTTPClient import NVPair

connectionDefaults = HTTPPluginControl.getConnectionDefaults()
httpUtilities = HTTPPluginControl.getHTTPUtilities()

from jcifs.ntlmssp import Type1Message, Type2Message, Type3Message
from jcifs.util import Base64

test1 = Test(1, "Request resource")
request1 = test1.wrap(HTTPRequest())

username="e461936"
password=""
domain="CORP.statestr.com"
workstation="CN-PC-HZ1239"

siteurl = "http://sstc"
url0 = "/sites/sitelist/ooa/peg/default.aspx"
url1 = "/sites/sitelist/ooa/peg/_layouts/1033/viewlsts.aspx?BaseType=0"

connectionDefaults.defaultHeaders = \
  ( NVPair('User-Agent', 'Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1; SV1; Avant Browser; InfoPath.1; .NET CLR 2.0.50727; .NET CLR 1.1.4322)'),
    NVPair('Accept-Encoding', 'gzip, deflate'),
    NVPair('Accept-Language', 'en-us'),
    NVPair("Connection", "Keep-Alive"))
#connectionDefaults.setProxyServer("127.0.0.1", 8888)

headers0 = \
  ( NVPair('Accept', 'image/gif, image/x-xbitmap, image/jpeg, image/pjpeg, application/x-shockwave-flash, application/vnd.ms-excel, application/vnd.ms-powerpoint, application/msword, application/x-silverlight, */*'),
    NVPair('Referer', siteurl + url0),
    NVPair("Connection", "Keep-Alive"), )

headers1 = \
  ( NVPair('Accept', 'image/gif, image/x-xbitmap, image/jpeg, image/pjpeg, application/x-shockwave-flash, application/vnd.ms-excel, application/vnd.ms-powerpoint, application/msword, application/x-silverlight, */*'),
    NVPair('Referer', siteurl + url1),
    NVPair("Connection", "Keep-Alive"), )

request0 = HTTPRequest(url=siteurl, headers=headers0)
request1 = HTTPRequest(url=siteurl, headers=headers1)

class TestRunner:
    def __init__(self):
        self.cnt = 0
        self.info = NTLMAuthenticationInfo(domain, workstation, username, password)
        request = request0
        result = request.GET(url0)

        result = NTLMAuthentication(result, request, self.info)
        self.cookies = CookieModule.listAllCookies( HTTPPluginControl.getThreadHTTPClientContext() )

    def __call__(self):
        pass
        for cookie in self.cookies:
            CookieModule.addCookie( cookie, HTTPPluginControl.getThreadHTTPClientContext() )
        request = request1

        result = request.GET(url1)
        result = NTLMAuthentication(result, request, self.info)
        self.cnt += 1
        print "self.cnt ", self.cnt

def NTLMAuthentication1(url, request, info, NTLMfield):
    token_type1 = info.token_type1()

    params = (NVPair("Authorization", "NTLM "+token_type1), )
    result = request.GET(url, None, params)
    NTLMfield = result.getHeader("WWW-Authenticate")
    return NTLMAuthentication2(url, request, info, NTLMfield)

def NTLMAuthentication2(url, request, info, NTLMfield):
    if NTLMfield.startswith("Negotiate"):
        token_type2 = NTLMfield[len("Negotiate "):]
    else:
        token_type2 = NTLMfield[5:]

    token_type3 = info.token_type3(token_type2)
    params = (NVPair("Cookie", "WSS_KeepSessionAuthenticated=80"),
              NVPair("Authorization", "NTLM " + token_type3), )
    result = request.GET(url, None, params)
    return result

# this function validate request and its result to see if the NTLM authentication is required
def NTLMAuthentication(lastResult, request, info):
    # get last http request's url
    url = lastResult.getEffectiveURI().toString()[len(request.getUrl()):]

    # The result is ask for authentication
    if lastResult.statusCode != 401 and lastResult.statusCode != 407:
        return lastResult

    NTLMfield = lastResult.getHeader("WWW-Authenticate")
    if NTLMfield == None:
        return lastResult

    # check it is the first shakehands
    if NTLMfield == "Negotiate, NTLM" or NTLMfield == "NTLM":
        return NTLMAuthentication1(url, request, info, NTLMfield)

    # check it is the second shakehands
    elif len(NTLMfield) > 4 and NTLMfield[:4] == "NTLM":
        return NTLMAuthentication2(url, request, info, NTLMfield)

    else:
        return lastResult

class NTLMAuthenticationInfo:
    def __init__(self, domain, host, user, passwd):
        self.domain = domain
        self.host = host
        self.user = user
        self.passwd = passwd

    def token_type1(self):
        msg = Type1Message(Type1Message.getDefaultFlags(), self.domain, self.host)
        return Base64.encode(msg.toByteArray())

    def token_type3(self, token_type2):
        msg2 = Type2Message(Base64.decode(token_type2))

#if jcifs 1.3.7 using msg3 = Type3Message(msg2, self.passwd, self.domain, self.user, self.host)
        msg3 = Type3Message(msg2, self.passwd, self.domain, self.user, self.host)
        return Base64.encode(msg3.toByteArray())

# Utility method that writes the given string to a uniquely named file
# using a FilenameFactory.
def writeToFile(text):
    filename = grinder.getFilenameFactory().createFilename(
        "page", "-%d.html" % grinder.runNumber)

    file = open(filename, "w")
    print >> file, text
    file.close()