Skip to content

Commit 5739035

Browse files
gpotter2Copilot
andauthored
Improve AutoArgparse & rework extensions (#4806)
* Improve AutoArgparse logic * Rework extensions loading * Update scapy/config.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 5744d95 commit 5739035

File tree

3 files changed

+184
-87
lines changed

3 files changed

+184
-87
lines changed

scapy/config.py

Lines changed: 83 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,11 @@
77
Implementation of the configuration object.
88
"""
99

10-
1110
import atexit
1211
import copy
1312
import functools
1413
import os
14+
import pathlib
1515
import re
1616
import socket
1717
import sys
@@ -538,7 +538,7 @@ def __repr__(self):
538538

539539

540540
class ScapyExt:
541-
__slots__ = ["specs", "name", "version"]
541+
__slots__ = ["specs", "name", "version", "bash_completions"]
542542

543543
class MODE(Enum):
544544
LAYERS = "layers"
@@ -554,6 +554,7 @@ class ScapyExtSpec:
554554

555555
def __init__(self):
556556
self.specs: Dict[str, 'ScapyExt.ScapyExtSpec'] = {}
557+
self.bash_completions = {}
557558

558559
def config(self, name, version):
559560
self.name = name
@@ -576,6 +577,9 @@ def register(self, name, mode, path, default=None):
576577
spec.default = bool(importlib.util.find_spec(spec.fullname))
577578
self.specs[fullname] = spec
578579

580+
def register_bashcompletion(self, script: pathlib.Path):
581+
self.bash_completions[script.name] = script
582+
579583
def __repr__(self):
580584
return "<ScapyExt %s %s (%s specs)>" % (
581585
self.name,
@@ -585,18 +589,19 @@ def __repr__(self):
585589

586590

587591
class ExtsManager(importlib.abc.MetaPathFinder):
588-
__slots__ = ["exts", "_loaded", "all_specs"]
592+
__slots__ = ["exts", "all_specs"]
589593

590-
SCAPY_PLUGIN_CLASSIFIER = 'Framework :: Scapy'
591-
GPLV2_CLASSIFIERS = [
592-
'License :: OSI Approved :: GNU General Public License v2 (GPLv2)',
593-
'License :: OSI Approved :: GNU General Public License v2 or later (GPLv2+)',
594+
GPLV2_LICENCES = [
595+
"GPL-2.0-only",
596+
"GPL-2.0-or-later",
594597
]
595598

596599
def __init__(self):
597600
self.exts: List[ScapyExt] = []
598601
self.all_specs: Dict[str, ScapyExt.ScapyExtSpec] = {}
599-
self._loaded = []
602+
# Add to meta_path as we are an import provider
603+
if self not in sys.meta_path:
604+
sys.meta_path.append(self)
600605

601606
def find_spec(self, fullname, path, target=None):
602607
if fullname in self.all_specs:
@@ -606,81 +611,87 @@ def invalidate_caches(self):
606611
pass
607612

608613
def _register_spec(self, spec):
614+
# Register to known specs
609615
self.all_specs[spec.fullname] = spec
616+
617+
# If default=True, inject it in the currently loaded modules
610618
if spec.default:
611619
loader = importlib.util.LazyLoader(spec.spec.loader)
612620
spec.spec.loader = loader
613621
module = importlib.util.module_from_spec(spec.spec)
614622
sys.modules[spec.fullname] = module
615623
loader.exec_module(module)
616624

617-
def load(self):
625+
def load(self, extension: str):
626+
"""
627+
Load a scapy extension.
628+
629+
:param extension: the name of the extension, as installed.
630+
"""
618631
try:
619632
import importlib.metadata
620633
except ImportError:
634+
raise ImportError("Cannot import importlib.metadata ! Upgrade Python.")
635+
636+
# Get extension distribution
637+
distr = importlib.metadata.distribution(extension)
638+
639+
# Check the classifiers
640+
if distr.metadata.get('License-Expression', None) not in self.GPLV2_LICENCES:
641+
log_loading.warning(
642+
"'%s' has no GPLv2 classifier therefore cannot be loaded." % extension
643+
)
621644
return
622-
for distr in importlib.metadata.distributions():
623-
if any(
624-
v == self.SCAPY_PLUGIN_CLASSIFIER
625-
for k, v in distr.metadata.items() if k == 'Classifier'
626-
):
627-
try:
628-
# Python 3.13 raises an internal warning when calling this
629-
with warnings.catch_warnings():
630-
warnings.filterwarnings("ignore", category=DeprecationWarning)
631-
pkg = next(
632-
k
633-
for k, v in
634-
importlib.metadata.packages_distributions().items()
635-
if distr.name in v
636-
)
637-
except KeyError:
638-
pkg = distr.name
639-
if pkg in self._loaded:
640-
continue
641-
if not any(
642-
v in self.GPLV2_CLASSIFIERS
643-
for k, v in distr.metadata.items() if k == 'Classifier'
644-
):
645-
log_loading.warning(
646-
"'%s' has no GPLv2 classifier therefore cannot be loaded." % pkg # noqa: E501
647-
)
648-
continue
649-
self._loaded.append(pkg)
650-
ext = ScapyExt()
651-
try:
652-
scapy_ext = importlib.import_module(pkg)
653-
except Exception as ex:
654-
log_loading.warning(
655-
"'%s' failed during import with %s" % (
656-
pkg,
657-
ex
658-
)
659-
)
660-
continue
661-
try:
662-
scapy_ext_func = scapy_ext.scapy_ext
663-
except AttributeError:
664-
log_loading.info(
665-
"'%s' included the Scapy Framework specifier "
666-
"but did not include a scapy_ext" % pkg
667-
)
668-
continue
669-
try:
670-
scapy_ext_func(ext)
671-
except Exception as ex:
672-
log_loading.warning(
673-
"'%s' failed during initialization with %s" % (
674-
pkg,
675-
ex
676-
)
645+
646+
# Create the extension
647+
ext = ScapyExt()
648+
649+
# Get the top-level declared "import packages"
650+
# HACK: not available nicely in importlib :/
651+
packages = distr.read_text("top_level.txt").split()
652+
653+
for package in packages:
654+
scapy_ext = importlib.import_module(package)
655+
656+
# We initialize the plugin by calling it's 'scapy_ext' function
657+
try:
658+
scapy_ext_func = scapy_ext.scapy_ext
659+
except AttributeError:
660+
log_loading.warning(
661+
"'%s' does not look like a Scapy plugin !" % extension
662+
)
663+
return
664+
try:
665+
scapy_ext_func(ext)
666+
except Exception as ex:
667+
log_loading.warning(
668+
"'%s' failed during initialization with %s" % (
669+
extension,
670+
ex
677671
)
678-
continue
679-
for spec in ext.specs.values():
680-
self._register_spec(spec)
681-
self.exts.append(ext)
682-
if self not in sys.meta_path:
683-
sys.meta_path.append(self)
672+
)
673+
return
674+
675+
# Register all the specs provided by this extension
676+
for spec in ext.specs.values():
677+
self._register_spec(spec)
678+
679+
# Add to the extension list
680+
self.exts.append(ext)
681+
682+
# If there are bash autocompletions, add them
683+
if ext.bash_completions:
684+
from scapy.main import _add_bash_autocompletion
685+
686+
for name, script in ext.bash_completions.items():
687+
_add_bash_autocompletion(name, script)
688+
689+
def loadall(self) -> None:
690+
"""
691+
Load all extensions registered in conf.
692+
"""
693+
for extension in conf.load_extensions:
694+
self.load(extension)
684695

685696
def __repr__(self):
686697
from scapy.utils import pretty_list
@@ -1033,6 +1044,8 @@ class Conf(ConfClass):
10331044
#: netcache holds time-based caches for net operations
10341045
netcache: NetCache = NetCache()
10351046
geoip_city = None
1047+
#: Scapy extensions that are loaded automatically on load
1048+
load_extensions: List[str] = []
10361049
# can, tls, http and a few others are not loaded by default
10371050
load_layers: List[str] = [
10381051
'bluetooth',
@@ -1170,10 +1183,6 @@ def __getattribute__(self, attr):
11701183

11711184
conf = Conf() # type: Conf
11721185

1173-
# Python 3.8 Only
1174-
if sys.version_info >= (3, 8):
1175-
conf.exts.load()
1176-
11771186

11781187
def crypto_validator(func):
11791188
# type: (DecoratorCallable) -> DecoratorCallable

scapy/main.py

Lines changed: 68 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,20 +9,22 @@
99

1010

1111
import builtins
12-
import pathlib
13-
import sys
14-
import os
15-
import getopt
1612
import code
17-
import gzip
13+
import getopt
1814
import glob
15+
import gzip
1916
import importlib
2017
import io
21-
from itertools import zip_longest
2218
import logging
19+
import os
20+
import pathlib
2321
import pickle
22+
import shutil
23+
import sys
2424
import types
2525
import warnings
26+
27+
from itertools import zip_longest
2628
from random import choice
2729

2830
# Never add any global import, in main.py, that would trigger a
@@ -101,6 +103,15 @@ def _probe_cache_folder(*cf):
101103
)
102104

103105

106+
def _probe_share_folder(*cf):
107+
# type: (str) -> Optional[pathlib.Path]
108+
return _probe_xdg_folder(
109+
"XDG_DATA_HOME",
110+
os.path.join(os.path.expanduser("~"), ".local", "share"),
111+
*cf
112+
)
113+
114+
104115
def _check_perms(file: Union[pathlib.Path, str]) -> None:
105116
"""
106117
Checks that the permissions of a file are properly user-specific, if sudo is used.
@@ -203,6 +214,22 @@ def _validate_local(k):
203214
DEFAULT_PRESTART_FILE = None
204215
DEFAULT_STARTUP_FILE = None
205216

217+
# https://github.com/scop/bash-completion/blob/main/README.md#faq
218+
if "BASH_COMPLETION_USER_DIR" in os.environ:
219+
BASH_COMPLETION_USER_DIR: Optional[pathlib.Path] = pathlib.Path(
220+
os.environ["BASH_COMPLETION_USER_DIR"]
221+
)
222+
else:
223+
BASH_COMPLETION_USER_DIR = _probe_share_folder("bash-completion")
224+
225+
if BASH_COMPLETION_USER_DIR:
226+
BASH_COMPLETION_FOLDER: Optional[pathlib.Path] = (
227+
BASH_COMPLETION_USER_DIR / "completions"
228+
)
229+
else:
230+
BASH_COMPLETION_FOLDER = None
231+
232+
206233
# Default scapy prestart.py config file
207234

208235
DEFAULT_PRESTART = """
@@ -219,6 +246,12 @@ def _validate_local(k):
219246
# disable INFO: tags related to dependencies missing
220247
# log_loading.setLevel(logging.WARNING)
221248
249+
# extensions to load by default
250+
conf.load_extensions = [
251+
# "scapy-red",
252+
# "scapy-rpc",
253+
]
254+
222255
# force-use libpcap
223256
# conf.use_pcap = True
224257
""".strip()
@@ -237,6 +270,31 @@ def _usage():
237270
sys.exit(0)
238271

239272

273+
def _add_bash_autocompletion(fname: str, script: pathlib.Path) -> None:
274+
"""
275+
Util function used most notably in setup.py to add a bash autocompletion script.
276+
"""
277+
try:
278+
if BASH_COMPLETION_FOLDER is None:
279+
raise OSError()
280+
281+
# If already defined, exit.
282+
dest = BASH_COMPLETION_FOLDER / fname
283+
if dest.exists():
284+
return
285+
286+
# Check that bash autocompletion folder exists
287+
if not BASH_COMPLETION_FOLDER.exists():
288+
BASH_COMPLETION_FOLDER.mkdir(parents=True, exist_ok=True)
289+
_check_perms(BASH_COMPLETION_FOLDER)
290+
291+
# Copy file
292+
shutil.copy(script, BASH_COMPLETION_FOLDER)
293+
except OSError:
294+
log_loading.warning("Bash autocompletion script could not be copied.",
295+
exc_info=True)
296+
297+
240298
######################
241299
# Extension system #
242300
######################
@@ -808,6 +866,10 @@ def interact(mydict=None,
808866
_locals=SESSION
809867
)
810868

869+
# Load extensions (Python 3.8 Only)
870+
if sys.version_info >= (3, 8):
871+
conf.exts.loadall()
872+
811873
if conf.fancy_banner:
812874
banner_text = get_fancy_banner()
813875
else:

0 commit comments

Comments
 (0)