This commit is contained in:
Jack Robison 2019-05-22 02:14:09 -04:00
parent 8457ef54e2
commit f5f3cc5f95
No known key found for this signature in database
GPG key ID: DF25C68FE0239BB2
2 changed files with 11 additions and 33 deletions

View file

@ -93,28 +93,23 @@ class UPnP:
) )
return None return None
async def get_port_mapping_by_index(self, index: int) -> Optional[Tuple[str, int, str, int, str, bool, str, int]]: async def get_port_mapping_by_index(self, index: int) -> Tuple[str, int, str, int, str, bool, str, int]:
try: return await self.gateway.commands.GetGenericPortMappingEntry(NewPortMappingIndex=index)
redirect = await self.gateway.commands.GetGenericPortMappingEntry(NewPortMappingIndex=index)
return redirect
except UPnPError:
return None
async def get_redirects(self) -> List[Tuple[str, int, str, int, str, bool, str, int]]: async def get_redirects(self) -> List[Tuple[str, int, str, int, str, bool, str, int]]:
redirects: List[Tuple[str, int, str, int, str, bool, str, int]] = [] redirects: List[Tuple[str, int, str, int, str, bool, str, int]] = []
cnt = 0 cnt = 0
redirect: Optional[Tuple[str, int, str, int, str, bool, str, int]] = None
try: try:
redirect = await self.gateway.commands.GetGenericPortMappingEntry(NewPortMappingIndex=cnt) redirect = await self.get_port_mapping_by_index(cnt)
except UPnPError: except UPnPError:
pass return redirects
while redirect is not None: while redirect is not None:
redirects.append(redirect) redirects.append(redirect)
cnt += 1 cnt += 1
try: try:
redirect = await self.gateway.commands.GetGenericPortMappingEntry(NewPortMappingIndex=cnt) redirect = await self.get_port_mapping_by_index(cnt)
except UPnPError: except UPnPError:
pass break
return redirects return redirects
async def get_specific_port_mapping(self, external_port: int, protocol: str) -> Tuple[int, str, bool, str, int]: async def get_specific_port_mapping(self, external_port: int, protocol: str) -> Tuple[int, str, bool, str, int]:
@ -144,22 +139,8 @@ class UPnP:
raise UPnPError("unsupported protocol: {}".format(protocol)) raise UPnPError("unsupported protocol: {}".format(protocol))
_internal_port = int(internal_port or port) _internal_port = int(internal_port or port)
requested_port = int(_internal_port) requested_port = int(_internal_port)
redirect_tups: List[Tuple[str, int, str, int, str, bool, str, int]] = []
cnt = 0
port = int(port) port = int(port)
redirect: Optional[Tuple[str, int, str, int, str, bool, str, int]] = None redirect_tups = await self.get_redirects()
try:
redirect = await self.gateway.commands.GetGenericPortMappingEntry(NewPortMappingIndex=cnt)
except UPnPError:
pass
while redirect is not None:
redirect_tups.append(redirect)
cnt += 1
try:
redirect = await self.gateway.commands.GetGenericPortMappingEntry(NewPortMappingIndex=cnt)
except UPnPError as err:
if "ArrayIndex" in str(err):
break
redirects: Dict[Tuple[int, str], Tuple[str, int, str]] = { redirects: Dict[Tuple[int, str], Tuple[str, int, str]] = {
(ext_port, proto): (int_host, int_port, desc) (ext_port, proto): (int_host, int_port, desc)
@ -171,11 +152,7 @@ class UPnP:
if int_host == self.lan_address and int_port == requested_port and desc == description: if int_host == self.lan_address and int_port == requested_port and desc == description:
return port return port
port += 1 port += 1
await self.gateway.commands.AddPortMapping( await self.add_port_mapping(port, protocol, _internal_port, self.lan_address, description)
NewRemoteHost='', NewExternalPort=port, NewProtocol=protocol,
NewInternalPort=_internal_port, NewInternalClient=self.lan_address,
NewEnabled=1, NewPortMappingDescription=description, NewLeaseDuration='0'
)
return port return port
# @cli # @cli

View file

@ -1,5 +1,6 @@
import unittest import unittest
from unittest import mock from unittest import mock
from aioupnp.upnp import UPnP
class mock_netifaces: class mock_netifaces:
@ -49,7 +50,7 @@ class TestParseInterfaces(unittest.TestCase):
def test_parse_interfaces(self): def test_parse_interfaces(self):
with mock.patch('aioupnp.interfaces.get_netifaces') as patch: with mock.patch('aioupnp.interfaces.get_netifaces') as patch:
patch.return_value = mock_netifaces patch.return_value = mock_netifaces
import aioupnp.interfaces
gateway, lan = aioupnp.interfaces.get_gateway_and_lan_addresses('test0') lan, gateway = UPnP.get_lan_and_gateway(interface_name='test0')
self.assertEqual(gateway, '192.168.1.1') self.assertEqual(gateway, '192.168.1.1')
self.assertEqual(lan, '192.168.1.2') self.assertEqual(lan, '192.168.1.2')