Imports System.DirectoryServices
Imports System.Web.Helpers
Imports RADAR

Public Class RDSHandler

    ' User -> Token that proves user has authenticated, but not yet proved
    ' herself with the 2. factor
    Private Shared authTokens As New Hashtable

    Private Shared userSessions As New Hashtable
    Private Shared sessionTimestamps As New Hashtable
    Private Shared encryptedChallangeResults As New Hashtable
    Private Shared userLaunchTimestamps As New Hashtable

    

    Private mPacket As RADIUSPacket
    Private mUsername As String
    Private mPassword As String

    ' RDS specific values 
    Private mIsAppLaunchRequest As Boolean
    Private mIsGatewayRequest As Boolean
    Private mUseSMSFactor As Boolean
    Private mUseEmailFactor As Boolean

    Public Sub New(packet As RADIUSPacket)
        mPacket = packet

        mUsername = mPacket.UserName
        mPassword = mPacket.UserPassword

        CleanUsername()

        For Each atts As RADIUSAttribute In mPacket.Attributes.GetAllAttributes(RadiusAttributeType.VendorSpecific)
            Dim value As String = atts.GetVendorSpecific.VendorValue.ToString

            Select Case UCase(value)
                Case "LAUNCH"
                    mIsAppLaunchRequest = True
                Case "TSGATEWAY"
                    mIsGatewayRequest = True
                Case "SMS"
                    mUseSMSFactor = True
                Case "EMAIL"
                    mUseEmailFactor = True
            End Select
        Next
    End Sub

    Private Sub CleanUsername()
        ' RD Gateway sends EXAMPLE\username
        ' RD Web sends example\username or - TODO - even example.com\username
        If Not mUsername = Nothing Then
            mUsername = mUsername.ToLower
        End If
    End Sub

    Public Sub ProcessRequest()
        If mIsAppLaunchRequest Then
            ProcessAppLaunchRequest()
        ElseIf mIsGatewayRequest Then
            ProcessGatewayRequest()
        Else
            ProcessAccessRequest()
        End If

    End Sub

    ' Process the RDS specific App Launch request.
    ' These requests are sent when an app is clicked in RD Web.
    '
    ' It's checked whether the session is still valid. In which case, a
    ' window is opened for the user, where we allow the user to connect
    ' through the gateway, an Accept-Access is returned and the RD Web
    ' launches the RDP client.
    '
    ' NOTE: Requests contain the session GUID in the password attribute
    ' of the packet.
    Public Sub ProcessAppLaunchRequest()
        RDSFactor.AccessLog(mPacket, "AppLaunchRequest")

        ' When the packet is an AppLaunchRequest the password attribute contains the session id!
        Dim packetSessionId = mPassword
        Dim storedSessionId = userSessions(mUsername)

        If storedSessionId = Nothing Then
            RDSFactor.AccessLog(mPacket, "User has no session. MUST re-authenticate!")
            mPacket.RejectAccessRequest()
            Exit Sub
        End If

        If Not storedSessionId = packetSessionId Then
            RDSFactor.AccessLog(mPacket, "Stored session id didn't match packet session id!")
            mPacket.RejectAccessRequest()
            Exit Sub
        End If

        If HasValidSession(mUsername) Then
            RDSFactor.AccessLog(mPacket, "Opening window")
            ' Pro-long user session
            sessionTimestamps(mUsername) = Now
            ' Open gateway connection window
            userLaunchTimestamps(mUsername) = Now
            mPacket.AcceptAccessRequest()
            Exit Sub
        Else
            RDSFactor.AccessLog(mPacket, "Session timed out -- User MUST re-authenticate")
            userSessions.Remove(mUsername)
            sessionTimestamps.Remove(mUsername)
            mPacket.RejectAccessRequest()
        End If

    End Sub

    Public Shared Function HasValidLaunchWindow(username) As Boolean
        Dim timestamp = userLaunchTimestamps(username)

        Dim secondsSinceLaunch = DateDiff(DateInterval.Second, timestamp, Now)
        If secondsSinceLaunch < RDSFactor.LaunchTimeOut Then
            Return True
        Else
            Return False
        End If
    End Function

    Public Shared Function HasValidSession(username) As Boolean
        Dim id = userSessions(username)
        Dim timestamp = sessionTimestamps(username)

        Dim minsSinceLastActivity = DateDiff(DateInterval.Minute, timestamp, Now)
        If minsSinceLastActivity < RDSFactor.SessionTimeOut Then
            Return True
        Else
            Return False
        End If
    End Function

    ' Process the request from the Network Policy Server in the RDS Gateway.
    ' These are sent when an RDP client tries to connect through the Gateway.
    '
    ' Accept-Access is returned when the user has a
    ' * valid session; and a
    ' * valid app launch window
    '
    ' The launch window is closed after this request.
    '
    ' TODO: Fix race-condition RD Web vs. Gateway. Don't start RDP client in RD Web 
    ' before ensuring App Launch request was successful 
    Public Sub ProcessGatewayRequest()
        RDSFactor.AccessLog(mPacket, "Gateway Request")

        Dim sessionId = userSessions(mUsername)
        Dim launchTimestamp = userLaunchTimestamps(mUsername)
        Dim attributes As New RADIUSAttributes

        If sessionId = Nothing Or launchTimestamp = Nothing Then
            RDSFactor.AccessLog(mPacket, "User's has no launch window. User must re-authenticate")
            mPacket.RejectAccessRequest()
            Exit Sub
        End If

        Dim hasProxyState = mPacket.Attributes.AttributeExists(RadiusAttributeType.ProxyState)
        If hasProxyState Then
            Dim proxyState = mPacket.Attributes.GetFirstAttribute(RadiusAttributeType.ProxyState)
            attributes.Add(proxyState)
        End If

        If HasValidLaunchWindow(mUsername) Then
            RDSFactor.AccessLog(mPacket, "Opening gateway launch window")
            mPacket.AcceptAccessRequest(attributes)
        Else
            RDSFactor.AccessLog(mPacket, "Gateway launch window has timed out!")
            mPacket.RejectAccessRequest()
        End If

        RDSFactor.AccessLog(mPacket, "Removing gateway launch window")
        userLaunchTimestamps.Remove(mUsername)
    End Sub

    Public Sub ProcessAccessRequest()
        Dim hasState = mPacket.Attributes.AttributeExists(RadiusAttributeType.State)
        If hasState Then
            ' An Access-Request with a state is pr. definition a challange response.
            ProcessChallengeResponse()
            Exit Sub
        End If

        RDSFactor.AccessLog(mPacket, "AccessRequest")
        Try
            Dim ldapResult = Authenticate()

            If RDSFactor.EnableOTP Then
                TwoFactorChallenge()
                Exit Sub
            Else
                Accept()
            End If
        Catch ex As Exception
            RDSFactor.AccessLog(mPacket, "Authentication failed. Sending reject. Error: " & ex.Message)
            mPacket.RejectAccessRequest()
        End Try
    End Sub

    Private Sub Accept()
        RDSFactor.AccessLog(mPacket, "AcceptAccessRequest")
        Dim sGUID As String = System.Guid.NewGuid.ToString()
        userSessions(mUsername) = sGUID
        sessionTimestamps(mUsername) = Now

        Dim attributes As New RADIUSAttributes
        Dim guidAttribute As New RADIUSAttribute(RadiusAttributeType.ReplyMessage, sGUID)

        attributes.Add(guidAttribute)
        mPacket.AcceptAccessRequest(attributes)
    End Sub

    Private Sub ProcessChallengeResponse()
        Dim authToken = mPacket.Attributes.GetFirstAttribute(RadiusAttributeType.State).ToString
        If Not authToken = authTokens(mUsername) Then
            Throw New Exception("User is trying to respond to challange without valid auth token")
        End If

        ' When the packet is an Challange-Response the password attr. contains the encrypted result
        Dim userEncryptedResult = mPassword
        Dim localEncryptedResult = encryptedChallangeResults(mUsername)

        If localEncryptedResult = userEncryptedResult Then
            RDSFactor.AccessLog(mPacket, "ChallengeResponse Success")
            encryptedChallangeResults.Remove(mUsername)
            authTokens.Remove(mUsername)
            Accept()
        Else
            RDSFactor.AccessLog(mPacket, "Wrong challange code!")
            mPacket.RejectAccessRequest()
        End If
    End Sub

    Private Sub TwoFactorChallenge()
        Dim challangeCode = RDSFactor.GenerateCode
        Dim authToken = System.Guid.NewGuid.ToString
        Dim clientIP = mPacket.EndPoint.Address.ToString
        Dim sharedSecret = RDSFactor.secrets(clientIP)

        RDSFactor.AccessLog(mPacket, "Access Challange Code: " & challangeCode)

        If sharedSecret = Nothing Then
            Throw New Exception("No shared secret for client:" & clientIP)
        End If

        authTokens(mUsername) = authToken
        Dim encryptedChallangeResult = Crypto.SHA256(mUsername & challangeCode & sharedSecret)
        encryptedChallangeResults(mUsername) = encryptedChallangeResult

        If mUseSMSFactor Then
            RDSFactor.AccessLog(mPacket, "TODO: Send SMS")
        End If

        If mUseEmailFactor Then
            RDSFactor.AccessLog(mPacket, "TODO: Send Email")
        End If

        Dim attributes As New RADIUSAttributes

        Dim replyMessageAttr As New RADIUSAttribute(RadiusAttributeType.ReplyMessage, "SMS Token")
        Dim stateAttr As New RADIUSAttribute(RadiusAttributeType.State, authToken)

        attributes.Add(replyMessageAttr)
        attributes.Add(stateAttr)

        mPacket.SendAccessChallange(attributes)
    End Sub

    Private Function Authenticate() As System.DirectoryServices.SearchResult
        Dim password As String = mPacket.UserPassword
        Dim ldapDomain As String = RDSFactor.LDAPDomain

        RDSFactor.AccessLog(mPacket, "Authenticating with LDAP: " & "LDAP://" & ldapDomain)
        Dim dirEntry As New DirectoryEntry("LDAP://" & ldapDomain, mUsername, password)

        Dim obj As Object = dirEntry.NativeObject
        Dim search As New DirectorySearcher(dirEntry)

        If InStr(mUsername, "@") > 0 Then
            search.Filter = "(userPrincipalName=" + mUsername + ")"
        Else
            search.Filter = "(SAMAccountName=" + Split(mUsername, "\")(1) + ")"
        End If

        search.PropertiesToLoad.Add("distinguishedName")
        If RDSFactor.EnableOTP = True Then
            search.PropertiesToLoad.Add(RDSFactor.ADField)
            search.PropertiesToLoad.Add(RDSFactor.ADMailField)
        End If

        Dim result = search.FindOne()

        If IsDBNull(result) Then
            RDSFactor.AccessLog(mPacket, "Failed to authenticate with Active Directory")
            Throw New MissingUser
        End If

        Return result
    End Function

    Private Function LdapGetNumber(result As SearchResult) As String
        Dim mobile = result.Properties(RDSFactor.ADField)(0)
        mobile = Replace(mobile, "+", "")
        If mobile.Trim.Length = 0 Then
            RDSFactor.AccessLog(mPacket, "Unable to find correct phone number for user " & mUsername)
        End If
        Return mobile
    End Function

    Private Function LdapGetEmail(result As SearchResult) As String
        Dim email = result.Properties(RDSFactor.ADMailField)(0)

        If InStr(email, "@") = 0 Then
            RDSFactor.AccessLog(mPacket, "Unable to find correct email for user " & mUsername)
        End If
        Return email
    End Function

    Public Shared Sub Cleanup()
        RDSFactor.AccessLog("TimerCleanUp")

        Dim users = New ArrayList(userSessions.Keys)
        For Each username In users
            If Not HasValidSession(username) Then
                userSessions.Remove(username)
                sessionTimestamps.Remove(username)
                userLaunchTimestamps.Remove(username)
                encryptedChallangeResults.Remove(username)
                authTokens.Remove(username)
            End If
        Next
    End Sub

End Class