1
0
mirror of https://github.com/calebstewart/pwncat.git synced 2024-11-27 02:44:14 +01:00
pwncat/tests/conftest.py
Mitul16 0632beb992
chore: Pre-merge tasks completed
Oops, I missed the import for PlatformError :p
2022-02-16 16:19:35 +05:30

103 lines
2.6 KiB
Python

#!/usr/bin/env python3
import os
import time
import random
import socket
import string
import dataclasses
from io import StringIO
import pytest
from Crypto.PublicKey import RSA
from pwncat.channel import ChannelError
PLATFORM_MAP = {"ubuntu": "linux", "centos": "linux", "windows": "windows"}
def connection_details_for(name):
"""Get connection details from environment for the given
host type name (e.g. ubuntu, centos, windows)"""
if name not in PLATFORM_MAP:
pytest.skip(f"{name} is not a known target")
if (
f"{name.upper()}_HOST" not in os.environ
or f"{name.upper()}_BIND_PORT" not in os.environ
):
pytest.skip(f"{name} not available")
return {
"platform": PLATFORM_MAP[name],
"host": os.environ[f"{name.upper()}_HOST"],
"port": int(os.environ[f"{name.upper()}_BIND_PORT"]),
"protocol": "connect",
}
@pytest.fixture(params=["ubuntu", "centos"])
def linux_details(request):
"""Get available connection details for linux hosts"""
return connection_details_for(request.param)
@pytest.fixture(params=["windows"])
def windows_details(request):
"""Get available connection details for windows hosts"""
return connection_details_for(request.param)
def session_for(request):
# Grab details for this target
details = connection_details_for(request.param)
# Check if there are manager arguments
manager_args = getattr(
request.node.get_closest_marker("manager_config"), "args", {}
)
if not manager_args:
manager_args = {}
if "config" not in manager_args:
manager_args["config"] = StringIO(
"""
set -g db "memory://"
"""
)
import pwncat.manager
with pwncat.manager.Manager(**manager_args) as manager:
for i in range(3):
try:
session = manager.create_session(**details)
yield session
break
except ChannelError:
# This seems to be because of the container setup, so we just add
# a little sleep in
time.sleep(2)
else:
raise Exception("failed to connect to container")
@pytest.fixture(params=["windows", "ubuntu", "centos"])
def session(request):
"""Start a session with any platform"""
yield from session_for(request)
@pytest.fixture(params=["windows"])
def windows(request):
"""Start a windows session"""
yield from session_for(request)
@pytest.fixture(params=["ubuntu", "centos"])
def linux(request):
"""Start a linux session"""
yield from session_for(request)