You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

113 lines
3.6 KiB

  1. #nullable enable
  2. using System;
  3. using System.Collections.Generic;
  4. using System.IO;
  5. using System.Linq;
  6. using System.Runtime.InteropServices;
  7. using System.Text;
  8. using System.Threading.Tasks;
  9. namespace IPA.AntiMalware.ComAPI
  10. {
  11. internal class AmsiFileStream : IAmsiStream, IDisposable
  12. {
  13. private readonly FileInfo file;
  14. private readonly IntPtr session;
  15. public AmsiFileStream(FileInfo file, IntPtr session)
  16. {
  17. this.file = file;
  18. this.session = session;
  19. }
  20. public unsafe void GetAttribute([In] AmsiAttribute attribute, [In] uint dataSize, [Out] byte* buffer, out uint writtenData)
  21. {
  22. switch (attribute)
  23. {
  24. case AmsiAttribute.AppName:
  25. writtenData = WriteWString(AmsiConstants.AppName, dataSize, buffer);
  26. return;
  27. case AmsiAttribute.Session:
  28. *(IntPtr*)buffer = session;
  29. writtenData = (uint)sizeof(IntPtr);
  30. return;
  31. case AmsiAttribute.ContentName:
  32. writtenData = WriteWString(file.FullName, dataSize, buffer);
  33. return;
  34. case AmsiAttribute.ContentSize:
  35. *(ulong*)buffer = (ulong)file.Length;
  36. writtenData = sizeof(ulong);
  37. return;
  38. default:
  39. throw new NotImplementedException(); // return e_notimpl
  40. }
  41. static unsafe uint WriteWString(string str, uint dataSize, byte* buffer)
  42. {
  43. fixed (char* name = str)
  44. {
  45. return (uint)Encoding.Unicode.GetBytes(name, str.Length, buffer, (int)dataSize);
  46. }
  47. }
  48. }
  49. private FileStream? stream;
  50. private bool disposedValue;
  51. private readonly byte[] readBuffer = new byte[1024];
  52. public unsafe void Read([In] ulong position, [In] uint dataSize, [Out] byte* buffer, out uint readSize)
  53. {
  54. stream ??= file.OpenRead();
  55. stream.Position = (long)position;
  56. var bytesToRead = dataSize;
  57. readSize = 0;
  58. while (bytesToRead > 0)
  59. {
  60. var bytesRead = stream.Read(readBuffer, 0, (int)Math.Min(readBuffer.Length, bytesToRead));
  61. if (bytesRead == 0)
  62. {
  63. break;
  64. }
  65. fixed (byte* readBufferPtr = readBuffer)
  66. {
  67. Buffer.MemoryCopy(readBufferPtr, buffer + readSize, dataSize - readSize, bytesRead);
  68. }
  69. bytesToRead -= (uint)bytesRead;
  70. readSize += (uint)bytesRead;
  71. }
  72. }
  73. protected virtual void Dispose(bool disposing)
  74. {
  75. if (!disposedValue)
  76. {
  77. if (disposing)
  78. {
  79. stream?.Dispose();
  80. }
  81. disposedValue = true;
  82. }
  83. }
  84. // This does not have unmanagd resources, so it doesn't need to exist
  85. // ~AmsiFileStream()
  86. // {
  87. // // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
  88. // Dispose(disposing: false);
  89. // }
  90. public void Dispose()
  91. {
  92. // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
  93. Dispose(disposing: true);
  94. GC.SuppressFinalize(this);
  95. }
  96. }
  97. }