EtherAssist.psm1

# Requires -Version 5.1

# Set progress preference to improve performance with web requests
$ProgressPreference = 'SilentlyContinue'

#region Internal Helper Functions
function Format-EAResponse {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNull()]
        [PSCustomObject]$ResponseObject,
        
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$InputText,
        
        [ValidateNotNullOrEmpty()]
        [string]$InputType = "Question",
        
        [switch]$OutputAsJson,
        [switch]$OutputAsObject,
        [switch]$MuteQuestion,
        [switch]$MuteDateTime,
        [switch]$MuteAnswer
    )
    
    if (-not $ResponseObject.answer) {
        throw "Response object does not contain an answer"
    }
    
    try {
        if ($OutputAsJson)
        {
            $resultObject = [PSCustomObject]@{
                $InputType = $InputText
                Answer = $ResponseObject.answer
                TimeOfQuery = Get-Date -Format 'o'
                Metadata = @{
                    Status = $ResponseObject.status
                    Model = $ResponseObject.model
                }
            }
            return $resultObject | ConvertTo-Json -Depth 10
        }
        elseif ($OutputAsObject)
        {
            $resultObject = [PSCustomObject]@{
                $InputType = $InputText
                Answer = $ResponseObject.answer
                TimeOfQuery = Get-Date
                Metadata = @{
                    Status = $ResponseObject.status
                    Model = $ResponseObject.model
                }
            }
            return $resultObject
        }
        else
        {
            $responseMessage = @()
            if (-not $MuteQuestion)
            {
                $responseMessage += "$InputType`: $InputText"
            }
            if (-not $MuteDateTime)
            {
                $responseMessage += "Date/Time: $(Get-Date)"
            }
            
            $answerText = $ResponseObject.answer
            if (-not $MuteAnswer)
            {
                $responseMessage += "Answer: $answerText"
            }
            else
            {
                $responseMessage += $answerText
            }
            
            return $responseMessage -join "`n"
        }
    }
    catch {
        throw "Error formatting response: $_"
    }
}

function Set-EAApiConfig {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$ApiKey,
        
        [Parameter(Mandatory = $false)]
        [ValidateNotNullOrEmpty()]
        [ValidateScript({
            if ($_ -match '^https?://[a-z0-9-]+(\.[a-z0-9-]+)*(\.[a-z]{2,})(:[0-9]+)?(/.*)?$') {
                $true
            }
            else {
                throw "The URL '$_' is not valid. Please provide a valid HTTP/HTTPS URL."
            }
        })]
        [string]$ApiUrl = "https://api.etherassist.ai"
    )
    try {
        $KeyFilePath = Join-Path $env:USERPROFILE "EtherAssistApiKey.xml"
        $secureApiKey = ConvertTo-SecureString $ApiKey -AsPlainText -Force
        $settings = @{
            ApiKey = $secureApiKey
            ApiUrl = $ApiUrl
        }
        $settings | Export-Clixml -Path $KeyFilePath
        Write-Verbose "API configuration saved successfully. Using URL: $ApiUrl"
    }
    catch {
        Write-Error "Error storing API Key: $_"
    }
}

function Get-EAApiConfig {
    [CmdletBinding()]
    param()
    
    $KeyFilePath = Join-Path $env:USERPROFILE "EtherAssistApiKey.xml"
    if (Test-Path $KeyFilePath) {
        try {
            $settings = Import-Clixml -Path $KeyFilePath
            $settings.ApiKey = [System.Runtime.InteropServices.Marshal]::PtrToStringAuto([System.Runtime.InteropServices.Marshal]::SecureStringToBSTR($settings.ApiKey))
            return $settings
        }
        catch {
            Write-Error "Error retrieving API Key: $_"
        }
    }
    else {
        Write-Error "API Key and URL are not set. Use Set-EAApiConfig to configure."
    }
}

function Invoke-EtherAssistApi {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$Endpoint,
        
        [Parameter(Mandatory = $true)]
        [ValidateNotNull()]
        [object]$Body,
        
        [Parameter()]
        [ValidateRange(1, 300)]
        [int]$TimeoutSec = 30,
        
        [Parameter()]
        [ValidateRange(1, 10)]
        [int]$MaxRetries = 3,
        
        [Parameter()]
        [switch]$EnableLogging
    )
    
    $settings = Get-EAApiConfig
    if (-not $settings) {
        Write-Error "API settings are not set. Use Set-EAApiConfig to set the key and URL."
        return
    }
    
    $uri = "$($settings.ApiUrl)$Endpoint"
    $headers = @{
        Authorization = "Bearer $($settings.ApiKey)"
        'Content-Type' = 'application/json'
    }
    
    if ($EnableLogging) {
        Write-Verbose "Request URI: $uri"
        Write-Verbose "Headers: $($headers | ConvertTo-Json)"
        Write-Verbose "Body: $($Body | ConvertTo-Json)"
    }
    
    $retryCount = 0
    $success = $false
    
    while (-not $success -and $retryCount -lt $MaxRetries) {
        try {
            $response = Invoke-RestMethod -Uri $uri -Method Post -Headers $headers -Body ($Body | ConvertTo-Json) -TimeoutSec $TimeoutSec
            
            # Standardize the response format
            if ($response) {
                $standardResponse = @{
                    answer = $response.answer
                    status = $response.status
                    model = $response.model
                    success = $true
                }
                $success = $true
                return [PSCustomObject]$standardResponse
            }
        }
        catch {
            $retryCount++
            if ($retryCount -eq $MaxRetries) {
                Write-Error "API request failed after $MaxRetries attempts: $_"
                throw
            }
            Start-Sleep -Seconds ($retryCount * 2)
        }
    }
}

function Send-EARequest {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$Question,
        [switch]$UseAdvancedModel,
        [switch]$MuteQuestion,
        [switch]$MuteDateTime,
        [switch]$MuteAnswer,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $Question
        useAdvancedModel = $UseAdvancedModel.IsPresent
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/question" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $Question -InputType "Question" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject `
            -MuteQuestion:$MuteQuestion -MuteDateTime:$MuteDateTime -MuteAnswer:$MuteAnswer
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Invoke-EAQuery {
    [CmdletBinding()]
    param (
        [switch]$UseAdvancedModel,
        [switch]$MuteQuestion,
        [switch]$MuteDateTime,
        [switch]$MuteAnswer
    )
    
    try {
        $question = Read-Host "Please enter the question for EtherAssist"
        Send-EARequest -Question $question -UseAdvancedModel:$UseAdvancedModel -MuteQuestion:$MuteQuestion -MuteDateTime:$MuteDateTime -MuteAnswer:$MuteAnswer
    }
    catch {
        Write-Error "An error occurred while invoking EtherAssist query: $_"
    }
}

function Send-EAWebSearch {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$Query,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject,
        [switch]$MuteQuestion,
        [switch]$MuteDateTime,
        [switch]$MuteAnswer
    )
    
    $body = @{
        question = $Query
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/web-search" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $Query -InputType "SearchQuery" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject `
            -MuteQuestion:$MuteQuestion -MuteDateTime:$MuteDateTime -MuteAnswer:$MuteAnswer
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Convert-EABatchToPs {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$BatchCode,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $BatchCode
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/gen/convert-batch-to-ps" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $BatchCode -InputType "BatchScript" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function New-EAKnowledgeArticle {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$Topic,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject,
        [switch]$MuteQuestion,
        [switch]$MuteDateTime,
        [switch]$MuteAnswer
    )
    
    $body = @{
        question = $Topic
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/gen/create-knowledge-article" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $Topic -InputType "Topic" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject `
            -MuteQuestion:$MuteQuestion -MuteDateTime:$MuteDateTime -MuteAnswer:$MuteAnswer
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Set-EAApiAdvancedConfig {
    [CmdletBinding()]
    param (
        [Parameter()]
        [ValidateRange(1, 300)]
        [int]$DefaultTimeout = 30,
        
        [Parameter()]
        [ValidateRange(1, 10)]
        [int]$MaxRetries = 3,
        
        [Parameter()]
        [switch]$EnableLogging,
        
        [Parameter()]
        [ValidateScript({
            if ([string]::IsNullOrEmpty($_) -or $_ -match '^https?://[a-z0-9-]+(\.[a-z0-9-]+)*(\.[a-z]{2,})(:[0-9]+)?(/.*)?$') {
                $true
            }
            else {
                throw "The proxy URL '$_' is not valid. Please provide a valid HTTP/HTTPS URL or leave empty."
            }
        })]
        [string]$ProxyUrl
    )
    
    try {
        $ConfigFilePath = Join-Path $env:USERPROFILE "EtherAssistConfig.xml"
        
        $config = @{
            DefaultTimeout = $DefaultTimeout
            MaxRetries = $MaxRetries
            EnableLogging = $EnableLogging.IsPresent
            ProxyUrl = $ProxyUrl
            LastModified = Get-Date
        }
        
        # Encrypt sensitive data
        if ($ProxyUrl) {
            $secureProxyUrl = ConvertTo-SecureString $ProxyUrl -AsPlainText -Force
            $config.ProxyUrl = $secureProxyUrl
        }
        
        $config | Export-Clixml -Path $ConfigFilePath
        Write-Verbose "Advanced configuration saved successfully"
    }
    catch {
        Write-Error "Error saving advanced configuration: $_"
    }
}

function Get-EAApiAdvancedConfig {
    [CmdletBinding()]
    param()
    
    $ConfigFilePath = Join-Path $env:USERPROFILE "EtherAssistConfig.xml"
    
    if (Test-Path $ConfigFilePath) {
        try {
            $config = Import-Clixml -Path $ConfigFilePath
            
            # Decrypt sensitive data
            if ($config.ProxyUrl -is [System.Security.SecureString]) {
                $config.ProxyUrl = [System.Runtime.InteropServices.Marshal]::PtrToStringAuto(
                    [System.Runtime.InteropServices.Marshal]::SecureStringToBSTR($config.ProxyUrl)
                )
            }
            
            return [PSCustomObject]$config
        }
        catch {
            Write-Error "Error retrieving advanced configuration: $_"
        }
    }
    else {
        Write-Verbose "No advanced configuration found. Using defaults."
        return [PSCustomObject]@{
            DefaultTimeout = 30
            MaxRetries = 3
            EnableLogging = $false
            ProxyUrl = $null
            LastModified = $null
        }
    }
}

function Get-EATitle {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$Text,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $Text
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/utils/title" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $Text -InputType "Text" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EATextSummary {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$Text,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $Text
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/utils/Summarize" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $Text -InputType "Text" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EALogAnalysis {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$LogContent,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $LogContent
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/gen/log-analyze" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $LogContent -InputType "LogContent" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EAErrorCodeDescription {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$ErrorCode,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $ErrorCode
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/gen/errorcode" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $ErrorCode -InputType "ErrorCode" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Convert-EAVbsToPs {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$VbsCode,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $VbsCode
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/gen/convert-vbs-to-ps" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $VbsCode -InputType "VbsScript" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EAAppDescription {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$AppName,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $AppName
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/apps/app-description" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $AppName -InputType "AppName" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EAAppInstallerArgs {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$AppName,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $AppName
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/apps/installer-args" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $AppName -InputType "AppName" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EAMsixAnalysis {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$ManifestContent,
        [switch]$AppAttachAnalysis,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $ManifestContent
    }
    
    $endpoint = if ($AppAttachAnalysis) { "/apps/msix-app-attach-analysis" } else { "/apps/msix-analysis" }
    $responseObject = Invoke-EtherAssistApi -Endpoint $endpoint -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $ManifestContent -InputType "MsixManifest" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EAEntraThreatAnalysis {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$LogData,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $LogData
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/entra-threat-detection" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $LogData -InputType "EntraLogs" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

function Get-EAPcapAnalysis {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$PcapContent,
        [switch]$OutputAsJson,
        [switch]$OutputAsObject
    )
    
    $body = @{
        question = $PcapContent
    }
    
    $responseObject = Invoke-EtherAssistApi -Endpoint "/utils/pcap-analyze" -Body $body
    
    if ($responseObject.success -eq $true) {
        Format-EAResponse -ResponseObject $responseObject -InputText $PcapContent -InputType "PcapData" `
            -OutputAsJson:$OutputAsJson -OutputAsObject:$OutputAsObject
    }
    else {
        Write-Error "API request was not successful: $($responseObject.errorMessage)"
    }
}

Export-ModuleMember -Function Set-EAApiConfig, Get-EAApiConfig, Send-EARequest, Invoke-EAQuery, Send-EAWebSearch, Convert-EABatchToPs, New-EAKnowledgeArticle, Set-EAApiAdvancedConfig, Get-EAApiAdvancedConfig