137 lines
3.7 KiB
Nim
137 lines
3.7 KiB
Nim
import std/sequtils
|
|
import pkg/[
|
|
chronos,
|
|
chronicles
|
|
]
|
|
|
|
import ./tinyupnp
|
|
|
|
when (NimMajor, NimMinor) < (1, 4):
|
|
{.push raises: [Defect].}
|
|
else:
|
|
{.push raises: [].}
|
|
|
|
logScope: topics = "nat_mapping_manager"
|
|
|
|
type
|
|
PortMapping* = object
|
|
port*: int
|
|
protocol*: NatProtocolType
|
|
|
|
MappingManager* = ref object
|
|
mappings: seq[PortMapping]
|
|
upnpSession: TUpnpSession
|
|
agent*: string
|
|
mappingFrequency: Duration
|
|
refreshLoop: Future[void]
|
|
|
|
proc toUpnp(manager: MappingManager, pm: PortMapping): TUpnpPortMapping =
|
|
TUpnpPortMapping(
|
|
externalPort: pm.port,
|
|
internalPort: pm.port,
|
|
internalClient: manager.upnpSession.localIp,
|
|
protocol: pm.protocol,
|
|
description: manager.agent,
|
|
leaseDuration: manager.mappingFrequency
|
|
)
|
|
|
|
proc removeMappings*(manager: MappingManager, pms: seq[PortMapping]) {.async.} =
|
|
for mapping in pms:
|
|
if mapping in manager.mappings:
|
|
await manager.upnpSession.deletePortMapping(manager.toUpnp(mapping))
|
|
|
|
manager.mappings.keepItIf(it notin pms)
|
|
|
|
let allMappings = await manager.upnpSession.getAllMappings()
|
|
for mapping in pms:
|
|
let asUpnp = manager.toUpnp(mapping)
|
|
if allMappings.anyIt(it.same(asUpnp)):
|
|
# Nothing we can really do here except ringing the alarm
|
|
info "Can't delete upnp mapping!", mapping
|
|
|
|
proc removeAllMappings*(manager: MappingManager) {.async.} =
|
|
await manager.removeMappings(manager.mappings)
|
|
|
|
proc refreshMappings(manager: MappingManager) {.async.} =
|
|
for mapping in manager.mappings:
|
|
await manager.upnpSession.addPortMapping(manager.toUpnp(mapping))
|
|
|
|
proc refreshLoop(manager: MappingManager) {.async.} =
|
|
while true:
|
|
await sleepAsync(manager.mappingFrequency - 30.seconds)
|
|
try:
|
|
#TODO notify that our ip changed somehow
|
|
discard await manager.upnpSession.check()
|
|
await manager.refreshMappings()
|
|
except CatchableError as exc:
|
|
warn "Failed to refresh mappings!", err=exc.msg
|
|
|
|
proc setup*(manager: MappingManager) {.async.} =
|
|
if isNil(manager.upnpSession):
|
|
manager.upnpSession = TUpnpSession.new()
|
|
await manager.upnpSession.setup()
|
|
|
|
proc stop*(manager: MappingManager) {.async.} =
|
|
if not isNil(manager.refreshLoop):
|
|
await manager.removeAllMappings()
|
|
manager.refreshLoop.cancel()
|
|
manager.refreshLoop = nil
|
|
|
|
proc publicIp*(manager: MappingManager): Future[IpAddress] {.async.} =
|
|
await manager.setup()
|
|
return manager.upnpSession.publicIp
|
|
|
|
proc addMappings*(manager: MappingManager, pms: seq[PortMapping]) {.async.} =
|
|
await manager.setup()
|
|
|
|
for pm in pms:
|
|
if pm notin manager.mappings:
|
|
manager.mappings.add(pm)
|
|
|
|
await manager.refreshMappings()
|
|
|
|
if manager.mappings.len > 0 and isNil(manager.refreshLoop):
|
|
manager.refreshLoop = refreshLoop(manager)
|
|
elif manager.mappings.len == 0:
|
|
await manager.stop()
|
|
|
|
proc addMapping*(manager: MappingManager, pm: PortMapping) {.async.} =
|
|
await manager.addMappings(@[pm])
|
|
|
|
proc setMappings*(manager: MappingManager, pms: seq[PortMapping]) {.async.} =
|
|
var
|
|
toAdd: seq[PortMapping]
|
|
toRemove: seq[PortMapping]
|
|
|
|
for pm in pms:
|
|
if pm notin manager.mappings:
|
|
toAdd.add(pm)
|
|
|
|
for pm in manager.mappings:
|
|
if pm notin pms:
|
|
toRemove.add(pm)
|
|
|
|
if toRemove.len > 0:
|
|
await manager.removeMappings(pms)
|
|
|
|
if toAdd.len > 0:
|
|
await manager.addMappings(pms)
|
|
|
|
proc new*(
|
|
T: type[MappingManager],
|
|
frequency = 20.minutes,
|
|
agent = "nim upnp"): T =
|
|
|
|
doAssert frequency > 50.seconds
|
|
T(
|
|
agent: agent,
|
|
mappingFrequency: frequency
|
|
)
|
|
|
|
when isMainModule:
|
|
let manager = MappingManager.new(frequency = 1.minutes)
|
|
waitFor manager.setMappings(@[PortMapping(port: 55551, protocol: Tcp)])
|
|
waitFor sleepAsync(3.minutes)
|
|
waitFor manager.removeAllMappings()
|
|
waitFor sleepAsync(2.minutes)
|