/*
 * Copyright (c) 1999 - 2003
 * NetGroup, Politecnico di Torino (Italy)
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 * 3. Neither the name of the Politecnico di Torino nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include <GlobalConst.h>

#include <stdarg.h>
#include <ntddk.h>
#include <ntiologc.h>
#include <ndis.h>
#include "debug.h"
#include "packet.h"

#include "win_bpf.h"

//-------------------------------------------------------------------

NTSTATUS
NPF_OpenDumpFile(POPEN_INSTANCE Open , PUNICODE_STRING fileName, BOOLEAN Append)
{
   NTSTATUS ntStatus;
   IO_STATUS_BLOCK IoStatus;
   OBJECT_ATTRIBUTES ObjectAttributes;
   PWCHAR PathPrefix;
   USHORT PathLen;
   UNICODE_STRING FullFileName;
   ULONG FullFileNameLength;
   PDEVICE_OBJECT fsdDevice;
   
    IF_LOUD(DbgPrint("NPF: OpenDumpFile.\n");)

   if(fileName->Buffer[0] == L'\\' &&
      fileName->Buffer[1] == L'?' &&
      fileName->Buffer[2] == L'?' &&
      fileName->Buffer[3] == L'\\'
   ){
      PathLen = 0;
   }
   else{
      PathPrefix = L"\\??\\";
      PathLen = 8;
   }
   
   // Insert the correct path prefix.
   FullFileNameLength = PathLen + fileName->MaximumLength;
   
   FullFileName.Buffer = ExAllocatePoolWithTag(NonPagedPool, 
      FullFileNameLength,
      '0DWA');
   
   if (FullFileName.Buffer == NULL) {
      ntStatus = STATUS_INSUFFICIENT_RESOURCES;
      return ntStatus;
   }
   
   FullFileName.Length = PathLen;
   FullFileName.MaximumLength = (USHORT)FullFileNameLength;
   
   if(PathLen)
      RtlMoveMemory (FullFileName.Buffer, PathPrefix, PathLen);
   
   RtlAppendUnicodeStringToString (&FullFileName, fileName);
   
   IF_LOUD(DbgPrint( "Packet: Attempting to open %wZ\n", &FullFileName);)
   
   InitializeObjectAttributes ( &ObjectAttributes,
      &FullFileName,
      OBJ_CASE_INSENSITIVE,
      NULL,
      NULL );
   
   // Create the dump file
   ntStatus = ZwCreateFile( &Open->DumpFileHandle,
      SYNCHRONIZE | FILE_WRITE_DATA,
      &ObjectAttributes,
      &IoStatus,
      NULL,
      FILE_ATTRIBUTE_NORMAL,
      FILE_SHARE_READ,
      (Append)?FILE_OPEN_IF:FILE_SUPERSEDE,
      FILE_SYNCHRONOUS_IO_NONALERT,
      NULL,
      0 );

    if ( !NT_SUCCESS( ntStatus ) )
    {
        IF_LOUD(DbgPrint("NPF: Error opening file %x\n", ntStatus);)
      
        ExFreePool(FullFileName.Buffer);
      Open->DumpFileHandle=NULL;
        ntStatus = STATUS_NO_SUCH_FILE;
        return ntStatus;
    }
   
   ExFreePool(FullFileName.Buffer);
   
   ntStatus = ObReferenceObjectByHandle(Open->DumpFileHandle,
      FILE_WRITE_ACCESS,
      *IoFileObjectType,
      KernelMode,
      &Open->DumpFileObject,
      0);
   
    if ( !NT_SUCCESS( ntStatus ) )
    {
        IF_LOUD(DbgPrint("NPF: Error creating file, status=%x\n", ntStatus);)
         
      ZwClose( Open->DumpFileHandle );
      Open->DumpFileHandle=NULL;
      
        ntStatus = STATUS_NO_SUCH_FILE;
        return ntStatus;
    }
   
    fsdDevice = IoGetRelatedDeviceObject(Open->DumpFileObject);

   IF_LOUD(DbgPrint("NPF: Dump: write file created succesfully, status=%d \n",ntStatus);)

   return ntStatus;
}   

//-------------------------------------------------------------------

NTSTATUS
NPF_StartDump(POPEN_INSTANCE Open)
{
   NTSTATUS ntStatus;
   struct packet_file_header hdr;
   IO_STATUS_BLOCK IoStatus;

    IF_LOUD(DbgPrint("NPF: StartDump.\n");)

   // Init the file header
   hdr.magic = TCPDUMP_MAGIC;
   hdr.version_major = PCAP_VERSION_MAJOR;
   hdr.version_minor = PCAP_VERSION_MINOR;
   hdr.thiszone = 0; /*Currently not set*/
   hdr.snaplen = 1514;
   hdr.sigfigs = 0;

   // Detect the medium type
   switch (Open->Medium){
      
   case NdisMediumWan:
      hdr.linktype = DLT_EN10MB;
      break;
      
   case NdisMedium802_3:
      hdr.linktype = DLT_EN10MB;
      break;
      
   case NdisMediumFddi:
      hdr.linktype = DLT_FDDI;
      break;
      
   case NdisMedium802_5:         
      hdr.linktype = DLT_IEEE802;   
      break;
      
   case NdisMediumArcnet878_2:
      hdr.linktype = DLT_ARCNET;
      break;
      
   case NdisMediumAtm:
      hdr.linktype = DLT_ATM_RFC1483;
      break;
      
   default:
      hdr.linktype = DLT_EN10MB;
   }

   // Write the header.
   // We can use ZwWriteFile because we are in the context of the application
   ntStatus = ZwWriteFile(Open->DumpFileHandle,
      NULL,
      NULL,
      NULL,
      &IoStatus,
      &hdr,
      sizeof(hdr),
      NULL,
      NULL );

   
    if ( !NT_SUCCESS( ntStatus ) )
    {
        IF_LOUD(DbgPrint("NPF: Error dumping file %x\n", ntStatus);)
      
      ZwClose( Open->DumpFileHandle );
      Open->DumpFileHandle=NULL;
      
        ntStatus = STATUS_NO_SUCH_FILE;
        return ntStatus;
    }

   Open->DumpOffset.QuadPart=24;
         
   ntStatus = PsCreateSystemThread(&Open->DumpThreadHandle,
      THREAD_ALL_ACCESS,
      (ACCESS_MASK)0L,
      0,
      0,
      NPF_DumpThread,
      Open);
   
    if ( !NT_SUCCESS( ntStatus ) )
    {
        IF_LOUD(DbgPrint("NPF: Error creating dump thread, status=%x\n", ntStatus);)
      
      ZwClose( Open->DumpFileHandle );
      Open->DumpFileHandle=NULL;

        return ntStatus;
    }  

   ntStatus = ObReferenceObjectByHandle(Open->DumpThreadHandle,
      THREAD_ALL_ACCESS,
      NULL,
      KernelMode,
      &Open->DumpThreadObject,
      0);

    if ( !NT_SUCCESS( ntStatus ) )
    {
        IF_LOUD(DbgPrint("NPF: Error creating dump thread, status=%x\n", ntStatus);)
      
      ObDereferenceObject(Open->DumpFileObject);
      ZwClose( Open->DumpFileHandle );
      Open->DumpFileHandle=NULL;

        return ntStatus;
    }  

   
   return ntStatus;
   
}

//-------------------------------------------------------------------
// Dump Thread
//-------------------------------------------------------------------

VOID NPF_DumpThread(POPEN_INSTANCE Open)
{
//   ULONG      FrozenNic;

    IF_LOUD(DbgPrint("NPF: In the work routine.  Parameter = 0x%0x\n",Open);)

   while(TRUE){

      // Wait until some packets arrive or the timeout expires
      NdisWaitEvent(&Open->DumpEvent, 5000);  

      IF_LOUD(DbgPrint("NPF: Worker Thread - event signalled\n");)
         
      if(Open->DumpLimitReached ||
         Open->Size==0){      // BufSize=0 means that this instance was closed, or that the buffer is too
                           // small for any capture. In both cases it is better to end the dump

         IF_LOUD(DbgPrint("NPF: Worker Thread - Exiting happily\n");)
         IF_LOUD(DbgPrint("Thread: Dumpoffset=%I64d\n",Open->DumpOffset.QuadPart);)

         PsTerminateSystemThread(STATUS_SUCCESS);
         return;
      }
      
      NdisResetEvent(&Open->DumpEvent);

      // Write the content of the buffer to the file
      if(NPF_SaveCurrentBuffer(Open) != STATUS_SUCCESS){
         PsTerminateSystemThread(STATUS_SUCCESS);
         return;
      }
   
   }

}

//-------------------------------------------------------------------

NTSTATUS NPF_SaveCurrentBuffer(POPEN_INSTANCE Open)
{
#if 0

   Thead=Open->Bhead;
   Ttail=Open->Btail;
   TLastByte=Open->BLastByte;
   
    IF_LOUD(DbgPrint("NPF: NPF_SaveCurrentBuffer.\n");)

   // Get the address of the buffer
   CurrBuff=Open->Buffer;
   //
   // Fill the application buffer
   //
   if( Ttail < Thead )
   {
      if(Open->MaxDumpBytes &&
         (UINT)Open->DumpOffset.QuadPart /*+ GetBuffOccupation(Open)*/ > Open->MaxDumpBytes)
      {
         // Size limit reached
         UINT PktLen;
         
         SizeToDump = 0;
         
         // Scan the buffer to detect the exact amount of data to save
         while(TRUE){
            PktLen = ((struct sf_pkthdr*)(CurrBuff + Thead + SizeToDump))->caplen + sizeof(struct sf_pkthdr);
            
            if((UINT)Open->DumpOffset.QuadPart + SizeToDump + PktLen > Open->MaxDumpBytes)
               break;
            
            SizeToDump += PktLen;
         }
         
      }
      else
         SizeToDump = TLastByte-Thead;
      
      lMdl=IoAllocateMdl(CurrBuff+Thead, SizeToDump, FALSE, FALSE, NULL);
      if (lMdl == NULL)
      {
         // No memory: stop dump
         IF_LOUD(DbgPrint("NPF: dump thread: Failed to allocate Mdl\n");)
         return STATUS_UNSUCCESSFUL;
      }
      
      MmBuildMdlForNonPagedPool(lMdl);
      
      // Write to disk
      NPF_WriteDumpFile(Open->DumpFileObject,
         &Open->DumpOffset,
         SizeToDump,
         lMdl,
         &IoStatus);
      
      IoFreeMdl(lMdl);
      
      if(!NT_SUCCESS(IoStatus.Status)){
         // Error
         return STATUS_UNSUCCESSFUL;
      }
      
      if(SizeToDump != TLastByte-Thead){
         // Size limit reached.
         Open->DumpLimitReached = TRUE;
   
         // Awake the application
         KeSetEvent(Open->ReadEvent,0,FALSE);

         return STATUS_UNSUCCESSFUL;
      }
      
      // Update the packet buffer
      Open->DumpOffset.QuadPart+=(TLastByte-Thead);
      Open->BLastByte=Ttail;
      Open->Bhead=0;
   }

   if( Ttail > Thead ){
      
      if(Open->MaxDumpBytes &&
         (UINT)Open->DumpOffset.QuadPart /* +GetBuffOccupation(Open)*/ > Open->MaxDumpBytes)
      {
         // Size limit reached
         UINT PktLen;
                  
         SizeToDump = 0;
         
         // Scan the buffer to detect the exact amount of data to save
         while(Thead + SizeToDump < Ttail){

            PktLen = ((struct sf_pkthdr*)(CurrBuff + Thead + SizeToDump))->caplen + sizeof(struct sf_pkthdr);
            
            if((UINT)Open->DumpOffset.QuadPart + SizeToDump + PktLen > Open->MaxDumpBytes)
               break;
            
            SizeToDump += PktLen;
         }
         
      }
      else
         SizeToDump = Ttail-Thead;
            
      lMdl=IoAllocateMdl(CurrBuff+Thead, SizeToDump, FALSE, FALSE, NULL);
      if (lMdl == NULL)
      {
         // No memory: stop dump
         IF_LOUD(DbgPrint("NPF: dump thread: Failed to allocate Mdl\n");)
         return STATUS_UNSUCCESSFUL;
      }
      
      MmBuildMdlForNonPagedPool(lMdl);
      
      // Write to disk
      NPF_WriteDumpFile(Open->DumpFileObject,
         &Open->DumpOffset,
         SizeToDump,
         lMdl,
         &IoStatus);
      
      IoFreeMdl(lMdl);
      
      if(!NT_SUCCESS(IoStatus.Status)){
         // Error
         return STATUS_UNSUCCESSFUL;
      }
      
      if(SizeToDump != Ttail-Thead){
         // Size limit reached.
         Open->DumpLimitReached = TRUE;

         // Awake the application
         KeSetEvent(Open->ReadEvent,0,FALSE);
         
         return STATUS_UNSUCCESSFUL;
      }
      
      // Update the packet buffer
      Open->DumpOffset.QuadPart+=(Ttail-Thead);         
      Open->Bhead=Ttail;
      
   }
#endif
   return STATUS_SUCCESS;
}

//-------------------------------------------------------------------

NTSTATUS NPF_CloseDumpFile(POPEN_INSTANCE Open){
#if 0
    IF_LOUD(DbgPrint("NPF: NPF_CloseDumpFile.\n");)
    IF_LOUD(DbgPrint("Dumpoffset=%d\n",Open->DumpOffset.QuadPart);)

DbgPrint("1\n");
   // Consistency check
   if(Open->DumpFileHandle == NULL)
      return STATUS_UNSUCCESSFUL;

DbgPrint("2\n");
   ZwClose( Open->DumpFileHandle );

   ObDereferenceObject(Open->DumpFileObject);
/*
   if(Open->DumpLimitReached == TRUE)
      // Limit already reached: don't save the rest of the buffer.
      return STATUS_SUCCESS;
*/
DbgPrint("3\n");

   NPF_OpenDumpFile(Open,&Open->DumpFileName, TRUE);

   // Flush the buffer to file 
   NPF_SaveCurrentBuffer(Open);

   // Close The file
   ObDereferenceObject(Open->DumpFileObject);
   ZwClose( Open->DumpFileHandle );
   
   Open->DumpFileHandle = NULL;

   ObDereferenceObject(Open->DumpFileObject);
#endif
   return STATUS_SUCCESS;
}

//-------------------------------------------------------------------

static NTSTATUS PacketDumpCompletion(PDEVICE_OBJECT DeviceObject,
                                PIRP Irp,
                                PVOID Context)
{

    // Copy the status information back into the "user" IOSB
    *Irp->UserIosb = Irp->IoStatus;
    
    // Wake up the mainline code
    KeSetEvent(Irp->UserEvent, 0, FALSE);
          
    return STATUS_MORE_PROCESSING_REQUIRED;
}

//-------------------------------------------------------------------

VOID NPF_WriteDumpFile(PFILE_OBJECT FileObject,
                             PLARGE_INTEGER Offset,
                        ULONG Length,
                        PMDL Mdl,
                        PIO_STATUS_BLOCK IoStatusBlock)
{
    PIRP irp;
    KEVENT event;
    PIO_STACK_LOCATION ioStackLocation;
    PDEVICE_OBJECT fsdDevice = IoGetRelatedDeviceObject(FileObject);
 
    // Set up the event we'll use
    KeInitializeEvent(&event, SynchronizationEvent, FALSE);
    
    // Allocate and build the IRP we'll be sending to the FSD
    irp = IoAllocateIrp(fsdDevice->StackSize, FALSE);

    if (!irp) {
        // Allocation failed, presumably due to memory allocation failure
        IoStatusBlock->Status = STATUS_INSUFFICIENT_RESOURCES;
        IoStatusBlock->Information = 0;

      return;
    }
    
    irp->MdlAddress = Mdl;
    irp->UserEvent = &event;
    irp->UserIosb = IoStatusBlock;
    irp->Tail.Overlay.Thread = PsGetCurrentThread();
    irp->Tail.Overlay.OriginalFileObject= FileObject;    
    irp->RequestorMode = KernelMode;
    
    // Indicate that this is a WRITE operation
    irp->Flags = IRP_WRITE_OPERATION;    
    
    // Set up the next I/O stack location
    ioStackLocation = IoGetNextIrpStackLocation(irp);
    ioStackLocation->MajorFunction = IRP_MJ_WRITE;
    ioStackLocation->MinorFunction = 0;
    ioStackLocation->DeviceObject = fsdDevice;
    ioStackLocation->FileObject = FileObject;
    IoSetCompletionRoutine(irp, PacketDumpCompletion, 0, TRUE, TRUE, TRUE);    
    ioStackLocation->Parameters.Write.Length = Length;    
    ioStackLocation->Parameters.Write.ByteOffset = *Offset;
    

    // Send it on.  Ignore the return code
    (void) IoCallDriver(fsdDevice, irp);
     
    // Wait for the I/O to complete.
    KeWaitForSingleObject(&event, Executive, KernelMode, TRUE, 0);

    // Free the IRP now that we are done with it
    IoFreeIrp(irp);

    return;

}