-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
139 lines (109 loc) · 3.77 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
# -*- coding: utf-8 -*-
from types import ModuleType
from pathlib import Path
import errno
import functools
import importlib.util
import os
import sys
try:
import ida_pro
except ImportError as exception:
raise RuntimeError("Not running in IDA Pro") from exception
@functools.cache
def get_script_path() -> Path:
"""Get path to this script file.
Returns:
Path to script file. The path is absolute and contains no symbolic-link
components. The return value is cached.
Raises:
FileNotFoundError: If the path does not exist.
NotADirectoryError: If a file is accessed as a directory.
RuntimeError: If an infinite path resolution loop is detected.
"""
return Path(__file__).resolve(strict=True).parent
@functools.cache
def get_venv_path() -> Path:
"""Get path to virtual environment directory. The directory will be created
by `ida-venv` if it doesn't already exist. However, if the path already
exists, it should point to a directory, otherwise an exception is thrown.
Returns:
Path to virtual environment directory. The path is absolute and contains
no symbolic-link components. The return value is cached.
Raises:
NotADirectoryError: If the virtual environment path exists, but it's not
a directory.
"""
path = get_script_path() / ".venv"
if path.exists() and not path.is_dir():
raise NotADirectoryError(errno.ENOTDIR, os.strerror(errno.ENOTDIR), str(path))
return path
@functools.cache
def get_requirements() -> list[str]:
"""Read and return list of required package names from requirements.txt. The
requirements file should be present in the same directory where this script
resides.
Returns:
List of required package names. The return value is cached.
Raises:
FileNotFoundError: If requirements.txt does not exist or is not a file.
"""
path = get_script_path() / "requirements.txt"
if not path.is_file():
raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), str(path))
with open(path, "r", encoding="utf-8") as fp:
return fp.read().splitlines()
@functools.cache
def import_ida_venv() -> ModuleType:
"""Dynamically import `ida-venv` to avoid having to install it in the IDA
Pro plug-in directory.
Returns:
A module object representing the dynamically loaded `ida-venv` instance.
The return value is cached.
Raises:
NotADirectoryError: If the `ida-venv` path exists, but it's not a directory.
ModuleNotFoundError: If the `ida-venv` module is not found where it is
expected to be.
"""
path = get_script_path() / "ida-venv"
if not path.is_dir():
raise NotADirectoryError(errno.ENOTDIR, os.strerror(errno.ENOTDIR), str(path))
sys.path.append(str(path))
ida_venv = importlib.import_module("ida_venv")
sys.modules["ida_venv"] = ida_venv
return ida_venv
def main() -> int:
#
# Get path to the currently executing script.
#
script_path = get_script_path()
print(f"Script at {script_path}")
#
# Get path to virtual environment.
#
venv_path = get_venv_path()
print(f"Virtual environment at {venv_path}")
#
# Get list of dependencies.
#
requirements = get_requirements()
print(f"Dependencies {requirements}")
#
# Import ida-venv.
#
ida_venv = import_ida_venv()
print(f"Loaded ida_venv {ida_venv}")
#
# Activate virtual environment and run env.py.
#
ida_venv.run_script_in_env(
script_path=script_path / "env.py",
venv_path=venv_path,
dependencies=requirements,
)
return os.EX_OK
if __name__ == "__main__":
if os.environ.get("EXIT"):
ida_pro.qexit(main())
else:
main()