1import select
2import socket
3
4import numpy as np
5
6from urh.dev.native.Device import Device
7from urh.util.Logger import logger
8
9
10class RTLSDRTCP(Device):
11    MAXDATASIZE = 65536
12    ENDIAN = "big"
13    RTL_TCP_CONSTS = ["NULL", "centerFreq", "sampleRate", "tunerGainMode", "tunerGain", "freqCorrection", "tunerIFGain",
14                      "testMode", "agcMode", "directSampling", "offsetTuning", "rtlXtalFreq", "tunerXtalFreq",
15                      "gainByIndex", "bandwidth", "biasTee"]
16
17    DATA_TYPE = np.uint8
18
19    @staticmethod
20    def receive_sync(data_connection, ctrl_connection, device_number: int, center_freq: int, sample_rate: int,
21                     bandwidth: int, gain: int, freq_correction: int, direct_sampling_mode: int, device_ip: str,
22                     port: int):
23        # connect and initialize rtl_tcp
24        sdr = RTLSDRTCP(center_freq, gain, sample_rate, bandwidth, device_number)
25        sdr.open(ctrl_connection, device_ip, port)
26        if sdr.socket_is_open:
27            sdr.device_number = device_number
28            sdr.set_parameter("centerFreq", int(center_freq), ctrl_connection)
29            sdr.set_parameter("sampleRate", int(sample_rate), ctrl_connection)
30            sdr.set_parameter("bandwidth", int(bandwidth), ctrl_connection)
31            sdr.set_parameter("freqCorrection", int(freq_correction), ctrl_connection)
32            sdr.set_parameter("directSampling", int(direct_sampling_mode), ctrl_connection)
33            # Gain has to be set last, otherwise it does not get considered by RTL-SDR
34            sdr.set_parameter("tunerGain", int(gain), ctrl_connection)
35            exit_requested = False
36
37            while not exit_requested:
38                while ctrl_connection.poll():
39                    result = sdr.process_command(ctrl_connection.recv(), ctrl_connection)
40                    if result == "stop":
41                        exit_requested = True
42                        break
43
44                if not exit_requested:
45                    data_connection.send_bytes(sdr.read_sync())
46
47            logger.debug("RTLSDRTCP: closing device")
48            sdr.close()
49        else:
50            ctrl_connection.send("Could not connect to rtl_tcp:404")
51        ctrl_connection.send("close:0")
52        data_connection.close()
53        ctrl_connection.close()
54
55    def process_command(self, command, ctrl_connection, is_tx=False):
56        logger.debug("RTLSDRTCP: {}".format(command))
57        if command == self.Command.STOP.name:
58            return self.Command.STOP
59
60        tag, value = command
61        if tag == self.Command.SET_FREQUENCY.name:
62            logger.info("RTLSDRTCP: Set center freq to {0}".format(int(value)))
63            return self.set_parameter("centerFreq", int(value), ctrl_connection)
64
65        elif tag == self.Command.SET_RF_GAIN.name:
66            logger.info("RTLSDRTCP: Set tuner gain to {0}".format(int(value)))
67            return self.set_parameter("tunerGain", int(value), ctrl_connection)
68
69        elif tag == self.Command.SET_IF_GAIN.name:
70            logger.info("RTLSDRTCP: Set if gain to {0}".format(int(value)))
71            return self.set_parameter("tunerIFGain", int(value), ctrl_connection)
72
73        elif tag == self.Command.SET_SAMPLE_RATE.name:
74            logger.info("RTLSDRTCP: Set sample_rate to {0}".format(int(value)))
75            return self.set_parameter("sampleRate", int(value), ctrl_connection)
76
77        elif tag == self.Command.SET_BANDWIDTH.name:
78            logger.info("RTLSDRTCP: Set bandwidth to {0}".format(int(value)))
79            return self.set_parameter("bandwidth", int(value), ctrl_connection)
80
81        elif tag == self.Command.SET_FREQUENCY_CORRECTION.name:
82            logger.info("RTLSDRTCP: Set ppm correction to {0}".format(int(value)))
83            return self.set_parameter("freqCorrection", int(value), ctrl_connection)
84
85        elif tag == self.Command.SET_DIRECT_SAMPLING_MODE.name:
86            logger.info("RTLSDRTCP: Set direct sampling mode to {0}".format(int(value)))
87            return self.set_parameter("directSampling", int(value), ctrl_connection)
88
89    def __init__(self, freq, gain, srate, bandwidth, device_number, resume_on_full_receive_buffer=False):
90        super().__init__(center_freq=freq, sample_rate=srate, bandwidth=bandwidth,
91                         gain=gain, if_gain=1, baseband_gain=1,
92                         resume_on_full_receive_buffer=resume_on_full_receive_buffer)
93
94        # default class parameters
95        self.receive_process_function = self.receive_sync
96        self.device_number = device_number
97        self.socket_is_open = False
98        self.success = 0
99
100    @property
101    def receive_process_arguments(self):
102        return self.child_data_conn, self.child_ctrl_conn, self.device_number, self.frequency, self.sample_rate, \
103               self.bandwidth, self.gain, self.freq_correction, self.direct_sampling_mode, self.device_ip, self.port
104
105    def open(self, ctrl_connection, hostname="127.0.0.1", port=1234):
106        if not self.socket_is_open:
107            try:
108                # Create socket and connect
109                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)
110                # self.sock.settimeout(1.0)  # Timeout 1s
111                self.sock.connect((hostname, port))
112            except Exception as e:
113                self.socket_is_open = False
114                logger.info("Could not connect to rtl_tcp at {0}:{1} ({2})".format(hostname, port, e))
115                ctrl_connection.send("Could not connect to rtl_tcp at {0} [{1}] ({2}):1".format(hostname, port, e))
116                return False
117
118            try:
119                # Receive rtl_tcp initial data
120                init_data = self.sock.recv(self.MAXDATASIZE)
121
122                if len(init_data) != 12:
123                    return False
124                if init_data[0:4] != b'RTL0':
125                    return False
126
127                # Extract tuner name
128                tuner_number = int.from_bytes(init_data[4:8], self.ENDIAN)
129                if tuner_number == 1:
130                    self.tuner = "E4000"
131                elif tuner_number == 2:
132                    self.tuner = "FC0012"
133                elif tuner_number == 3:
134                    self.tuner = "FC0013"
135                elif tuner_number == 4:
136                    self.tuner = "FC2580"
137                elif tuner_number == 5:
138                    self.tuner = "R820T"
139                elif tuner_number == 6:
140                    self.tuner = "R828D"
141                else:
142                    self.tuner = "Unknown"
143
144                # Extract IF and RF gain
145                self.if_gain = int.from_bytes(init_data[8:10], self.ENDIAN)
146                self.rf_gain = int.from_bytes(init_data[10:12], self.ENDIAN)
147
148                logger.info(
149                    "Connected to rtl_tcp at {0}:{1} (Tuner: {2}, RF-Gain: {3}, IF-Gain: {4})".format(hostname, port,
150                                                                                                      self.tuner,
151                                                                                                      self.rf_gain,
152                                                                                                      self.if_gain))
153                ctrl_connection.send(
154                    "Connected to rtl_tcp at {0}[{1}] (Tuner={2}, RF-Gain={3}, IF-Gain={4}):0".format(hostname, port,
155                                                                                                      self.tuner,
156                                                                                                      self.rf_gain,
157                                                                                                      self.if_gain))
158            except Exception as e:
159                self.socket_is_open = False
160                logger.info("This is not a valid rtl_tcp server at {0}:{1} ({2})".format(hostname, port, e))
161                return False
162
163            self.socket_is_open = True
164
165    def close(self):
166        if self.socket_is_open:
167            self.socket_is_open = False
168        return self.sock.close()
169
170    def set_parameter(self, param: str, value: int, ctrl_connection):  # returns error (True/False)
171        if self.socket_is_open:
172            msg = self.RTL_TCP_CONSTS.index(param).to_bytes(1, self.ENDIAN)  # Set param at bits 0-7
173            msg += value.to_bytes(4, self.ENDIAN)  # Set value at bits 8-39
174            try:
175                self.sock.sendall(msg)  # Send data to rtl_tcp
176            except OSError as e:
177                self.sock.close()
178                logger.info("Could not set parameter {0}:{1} ({2})".format(param, value, e))
179                ctrl_connection.send("Could not set parameter {0} {1} ({2}):1".format(param, value, e))
180                return True
181        return False
182
183    def read_sync(self):
184        s_read, _, _ = select.select([self.sock], [], [], .1)
185        if self.sock in s_read:
186            return self.sock.recv(self.MAXDATASIZE)
187        else:
188            return b''
189
190    @staticmethod
191    def bytes_to_iq(buffer):
192        return np.subtract(np.frombuffer(buffer, dtype=np.int8), 127).reshape((-1, 2), order="C")
193