Skip to content

Commit f63c9d8

Browse files
authored
Support mmap as source for Package (#81)
Based on suggestion from @qkaiser. This will make it easier to integrate with unblob (see onekey-sec/unblob#1244).
1 parent 2ea585a commit f63c9d8

File tree

2 files changed

+56
-5
lines changed

2 files changed

+56
-5
lines changed

src/pymsi/package.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import copy
22
import io
3+
import mmap
34
from pathlib import Path
45
from typing import Iterator, Optional, Union
56

@@ -18,13 +19,14 @@
1819

1920

2021
class Package:
21-
def __init__(self, path_or_bytesio: Union[Path, io.BytesIO]):
22-
if isinstance(path_or_bytesio, io.BytesIO):
23-
self.path = None
24-
self.file = path_or_bytesio
25-
else:
22+
# TODO: consider typing.BinaryIO
23+
def __init__(self, path_or_bytesio: Union[Path, io.BytesIO, mmap.mmap]):
24+
if isinstance(path_or_bytesio, Path):
2625
self.path = path_or_bytesio.resolve(True)
2726
self.file = self.path.open("rb")
27+
else:
28+
self.path = None
29+
self.file = path_or_bytesio
2830
self.tables = {}
2931
self.ole = None
3032
self.summary = None

tests/io.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import mmap
2+
from pathlib import Path
3+
4+
import pytest
5+
6+
import pymsi
7+
8+
9+
def read_package(path_or_bytesio):
10+
with pymsi.Package(path_or_bytesio) as package:
11+
msi = pymsi.Msi(package)
12+
return msi
13+
14+
15+
# Function to read a package using Path
16+
def read_package_path(file_path):
17+
path = Path(file_path)
18+
return read_package(path)
19+
20+
21+
# Function to read a package using with open
22+
def read_package_with_open(file_path):
23+
with open(file_path, "rb") as f:
24+
return read_package(f)
25+
26+
27+
# Function to read a package using mmap
28+
def read_package_mmap(file_path):
29+
with open(file_path, "r+b") as f:
30+
mm = mmap.mmap(f.fileno(), 0)
31+
return read_package(mm)
32+
33+
34+
# Test cases
35+
@pytest.mark.parametrize(
36+
"read_package_func", [read_package_path, read_package_with_open, read_package_mmap]
37+
)
38+
def test_read_package(read_package_func):
39+
test_file = "powertoys.msi"
40+
41+
msi = read_package_func(test_file)
42+
43+
size = (msi.package.ole.nb_sect * msi.package.ole.sector_size) + 512
44+
45+
print(size)
46+
47+
48+
if __name__ == "__main__":
49+
pytest.main()

0 commit comments

Comments
 (0)