2021-02-04 18:10:37 +01:00
from collections import namedtuple
2022-12-27 13:52:45 +01:00
import json
import logging
import socket
import re
import struct
import time
2020-12-30 21:58:39 +01:00
from typing import Tuple , Iterator
2021-02-04 18:10:37 +01:00
from urllib . parse import urlparse , parse_qs
2020-12-30 21:58:39 +01:00
from backend import Backend , Change
2022-12-27 13:52:45 +01:00
from protocol import PacketType , PKT_CHANGE_TYPES , change_from_packet , packet_from_change , send_packet , recv_packet
2020-12-30 21:58:39 +01:00
2021-02-18 17:41:51 +01:00
# Total number of reconnection tries
2022-12-27 13:52:45 +01:00
RECONNECT_TRIES = 5
2021-02-18 17:41:51 +01:00
# Delay in seconds between reconnections (initial)
2022-12-27 13:52:45 +01:00
RECONNECT_DELAY = 5
2021-02-18 17:41:51 +01:00
# Scale delay factor after each failure
2022-12-27 13:52:45 +01:00
RECONNECT_DELAY_BACKOFF = 1.5
2021-02-18 17:41:51 +01:00
2021-02-08 19:55:57 +01:00
HostPortInfo = namedtuple ( ' HostPortInfo ' , [ ' host ' , ' port ' , ' addrtype ' ] )
SocketURLInfo = namedtuple ( ' SocketURLInfo ' , [ ' target ' , ' proxytype ' , ' proxytarget ' ] )
2021-02-04 18:10:37 +01:00
2021-02-08 19:55:57 +01:00
# Network address type.
2022-12-27 13:52:45 +01:00
2021-02-04 18:10:37 +01:00
class AddrType :
IPv4 = 0
IPv6 = 1
NAME = 2
2021-02-08 19:55:57 +01:00
# Proxy type. Only SOCKS5 supported at the moment as this is sufficient for Tor.
2022-12-27 13:52:45 +01:00
2021-02-08 19:55:57 +01:00
class ProxyType :
DIRECT = 0
SOCKS5 = 1
2021-02-04 18:10:37 +01:00
2022-12-27 13:52:45 +01:00
2021-02-08 19:55:57 +01:00
def parse_host_port ( path : str ) - > HostPortInfo :
''' Parse a host:port pair. '''
2022-12-27 13:52:45 +01:00
if path . startswith ( ' [ ' ) : # bracketed IPv6 address
2021-02-08 19:55:57 +01:00
eidx = path . find ( ' ] ' )
2021-02-04 18:10:37 +01:00
if eidx == - 1 :
raise ValueError ( ' Unterminated bracketed host address. ' )
2021-02-08 19:55:57 +01:00
host = path [ 1 : eidx ]
2021-02-04 18:10:37 +01:00
addrtype = AddrType . IPv6
eidx + = 1
2021-02-08 19:55:57 +01:00
if eidx > = len ( path ) or path [ eidx ] != ' : ' :
2021-02-04 18:10:37 +01:00
raise ValueError ( ' Port number missing. ' )
eidx + = 1
else :
2021-02-08 19:55:57 +01:00
eidx = path . find ( ' : ' )
2021-02-04 18:10:37 +01:00
if eidx == - 1 :
raise ValueError ( ' Port number missing. ' )
2021-02-08 19:55:57 +01:00
host = path [ 0 : eidx ]
2022-12-27 13:52:45 +01:00
if re . match ( r ' \ d+ \ . \ d+ \ . \ d+ \ . \ d+$ ' , host ) : # matches IPv4 address format
2021-02-04 18:10:37 +01:00
addrtype = AddrType . IPv4
else :
addrtype = AddrType . NAME
eidx + = 1
try :
2021-02-08 19:55:57 +01:00
port = int ( path [ eidx : ] )
2021-02-04 18:10:37 +01:00
except ValueError :
raise ValueError ( ' Invalid port number ' )
2021-02-08 19:55:57 +01:00
return HostPortInfo ( host = host , port = port , addrtype = addrtype )
2022-12-27 13:52:45 +01:00
2021-02-08 19:55:57 +01:00
def parse_socket_url ( destination : str ) - > SocketURLInfo :
''' Parse a socket: URL to extract the information contained in it. '''
url = urlparse ( destination )
if url . scheme != ' socket ' :
raise ValueError ( ' Scheme for socket backend must be socket:... ' )
target = parse_host_port ( url . path )
proxytype = ProxyType . DIRECT
proxytarget = None
2021-02-04 18:10:37 +01:00
# parse query parameters
# reject unknown parameters (currently all of them)
qs = parse_qs ( url . query )
2021-02-08 19:55:57 +01:00
for ( key , values ) in qs . items ( ) :
2022-12-27 13:52:45 +01:00
if key == ' proxy ' : # proxy=socks5:127.0.0.1:9050
2021-02-08 19:55:57 +01:00
if len ( values ) != 1 :
raise ValueError ( ' Proxy can only have one value ' )
2021-02-04 18:10:37 +01:00
2021-02-08 19:55:57 +01:00
( ptype , ptarget ) = values [ 0 ] . split ( ' : ' , 1 )
if ptype != ' socks5 ' :
raise ValueError ( ' Unknown proxy type ' + ptype )
proxytype = ProxyType . SOCKS5
proxytarget = parse_host_port ( ptarget )
else :
raise ValueError ( ' Unknown query string parameter ' + key )
return SocketURLInfo ( target = target , proxytype = proxytype , proxytarget = proxytarget )
2021-02-04 18:10:37 +01:00
2022-12-27 13:52:45 +01:00
2020-12-30 21:58:39 +01:00
class SocketBackend ( Backend ) :
def __init__ ( self , destination : str , create : bool ) :
self . version = None
self . prev_version = None
self . destination = destination
2021-02-04 18:10:37 +01:00
self . url = parse_socket_url ( destination )
2021-02-18 17:41:51 +01:00
self . connect ( )
2021-02-08 19:55:57 +01:00
2021-02-18 17:41:51 +01:00
def connect ( self ) :
2021-02-08 19:55:57 +01:00
if self . url . proxytype == ProxyType . DIRECT :
if self . url . target . addrtype == AddrType . IPv6 :
self . sock = socket . socket ( socket . AF_INET6 , socket . SOCK_STREAM )
2022-12-27 13:52:45 +01:00
else : # TODO NAME is assumed to be IPv4 for now
2021-02-08 19:55:57 +01:00
self . sock = socket . socket ( socket . AF_INET , socket . SOCK_STREAM )
else :
2022-12-27 16:25:48 +01:00
assert self . url . proxytype == ProxyType . SOCKS5
2021-02-08 19:55:57 +01:00
import socks
self . sock = socks . socksocket ( )
self . sock . set_proxy ( socks . SOCKS5 , self . url . proxytarget . host , self . url . proxytarget . port )
2021-02-18 17:41:51 +01:00
logging . info ( ' Connecting to {} : {} (addrtype {} , proxytype {} , proxytarget {} )... ' . format (
2022-12-27 13:52:45 +01:00
self . url . target . host , self . url . target . port , self . url . target . addrtype ,
self . url . proxytype , self . url . proxytarget ) )
2021-02-08 19:55:57 +01:00
self . sock . connect ( ( self . url . target . host , self . url . target . port ) )
2021-02-18 17:41:51 +01:00
logging . info ( ' Connected to {} ' . format ( self . destination ) )
2020-12-30 21:58:39 +01:00
def _send_packet ( self , typ : int , payload : bytes ) - > None :
send_packet ( self . sock , typ , payload )
def _recv_packet ( self ) - > Tuple [ int , bytes ] :
return recv_packet ( self . sock )
def initialize ( self ) - > bool :
'''
Initialize socket backend by request current metadata from server .
'''
logging . info ( ' Initializing backend ' )
2021-02-18 17:41:51 +01:00
self . _request_metadata ( )
2020-12-30 21:58:39 +01:00
logging . info ( ' Initialized SocketBackend: protocol= {} , version= {} , prev_version= {} , version_count= {} ' . format (
self . protocol , self . version , self . prev_version , self . version_count
) )
return True
2021-02-18 17:41:51 +01:00
def _request_metadata ( self ) - > None :
self . _send_packet ( PacketType . REQ_METADATA , b ' ' )
( typ , payload ) = self . _recv_packet ( )
2022-12-27 16:25:48 +01:00
assert typ == PacketType . METADATA
2021-02-18 17:41:51 +01:00
self . protocol , self . version , self . prev_version , self . version_count = struct . unpack ( " !IIIQ " , payload )
2020-12-30 21:58:39 +01:00
def add_change ( self , entry : Change ) - > bool :
typ , payload = packet_from_change ( entry )
2021-02-18 17:41:51 +01:00
base_version = self . version
retry = 0
retry_delay = RECONNECT_DELAY
need_connect = False
2022-12-27 13:52:45 +01:00
while True : # Retry loop
2021-02-18 17:41:51 +01:00
try :
if need_connect :
self . connect ( )
# Request metadata, to know where we stand
self . _request_metadata ( )
if self . version == entry . version :
# If the current version at the server side matches the version of the
# entry, the packet was succesfully sent and processed and the error
# happened afterward. Nothing left to do.
return True
elif base_version == self . version :
# The other acceptable option is that the current version still matches
# that on the server side. Then we retry.
pass
else :
raise Exception ( ' Unexpected backup version {} after reconnect ' . format ( self . version ) )
self . _send_packet ( typ , payload )
# Wait for change to be acknowledged before continuing.
( typ , _ ) = self . _recv_packet ( )
2022-12-27 16:25:48 +01:00
assert typ == PacketType . ACK
2021-02-18 17:41:51 +01:00
except ( BrokenPipeError , OSError ) :
pass
else :
break
if retry == RECONNECT_TRIES :
logging . error ( ' Connection was lost while sending change (giving up after {} retries) ' . format ( retry ) )
raise IOError ( ' Connection was lost while sending change ' )
retry + = 1
logging . warning ( ' Connection was lost while sending change (retry {} of {} , will try again after {} seconds) ' . format ( retry , RECONNECT_TRIES , retry_delay ) )
time . sleep ( retry_delay )
retry_delay * = RECONNECT_DELAY_BACKOFF
need_connect = True
self . prev_version = self . version
self . version = entry . version
2020-12-30 21:58:39 +01:00
return True
def rewind ( self ) - > bool :
''' Rewind to previous version. '''
version = struct . pack ( " !I " , self . prev_version )
self . _send_packet ( PacketType . REWIND , version )
# Wait for change to be acknowledged before continuing.
( typ , _ ) = self . _recv_packet ( )
2022-12-27 16:25:48 +01:00
assert typ == PacketType . ACK
2020-12-30 21:58:39 +01:00
return True
def stream_changes ( self ) - > Iterator [ Change ] :
self . _send_packet ( PacketType . RESTORE , b ' ' )
version = - 1
while True :
( typ , payload ) = self . _recv_packet ( )
if typ in PKT_CHANGE_TYPES :
change = change_from_packet ( typ , payload )
version = change . version
yield change
elif typ == PacketType . DONE :
break
else :
raise ValueError ( " Unknown entry type {} " . format ( typ ) )
if version != self . version :
raise ValueError ( " Versions do not match up: restored version {} , backend version {} " . format ( version , self . version ) )
2022-12-27 16:25:48 +01:00
assert version == self . version
2020-12-30 21:58:39 +01:00
def compact ( self ) :
self . _send_packet ( PacketType . COMPACT , b ' ' )
( typ , payload ) = self . _recv_packet ( )
2022-12-27 16:25:48 +01:00
assert typ == PacketType . COMPACT_RES
2020-12-30 21:58:39 +01:00
return json . loads ( payload . decode ( ) )